Skip to content

Commit

Permalink
Refactor description handling
Browse files Browse the repository at this point in the history 
Description arguments will not work when specifying a tracker argument.

When automated tracker searching (ie: no tracker specified), discarding the found description will allow custom description use.
  • Loading branch information
@Audionut
Audionut committed Sep 29, 2024
1 parent 08db864 commit 04f4074
Show file tree
Hide file tree
Showing 3 changed files with 83 additions and 113 deletions.
183 changes: 71 additions & 112 deletions src/prep.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -174,6 +174,7 @@ async def update_metadata_from_tracker(self, tracker_name, tracker_instance, met
tracker_name,
tracker_instance.torrent_url,
tracker_instance.search_url,
meta,
file_name=search_term
)

Expand All @@ -198,30 +199,25 @@ async def update_metadata_from_tracker(self, tracker_name, tracker_instance, met
console.print(f"[green]{tracker_name} IMDb ID found: tt{meta['imdb']}[/green]")
if not meta['unattended']:
if await self.prompt_user_for_confirmation("Do you want to use this ID data from PTP?"):
meta['skip_gen_desc'] = True
found_match = True

# Retrieve PTP description and image list
ptp_desc, ptp_imagelist = await tracker_instance.get_ptp_description(ptp_torrent_id, meta, meta.get('is_disc', False))
meta['description'] = ptp_desc

if not meta.get('image_list'): # Only handle images if image_list is not already populated
valid_images = await self.check_images_concurrently(ptp_imagelist)
if valid_images:
meta['image_list'] = valid_images
await self.handle_image_list(meta, tracker_name)

meta['skip_gen_desc'] = True
console.print("[green]PTP images added to metadata.[/green]")

else:
found_match = False
meta['skip_gen_desc'] = True
meta['description'] = None

else:
found_match = True
ptp_desc, ptp_imagelist = await tracker_instance.get_ptp_description(ptp_torrent_id, meta, meta.get('is_disc', False))
meta['description'] = ptp_desc
meta['skip_gen_desc'] = True

if not meta.get('image_list'): # Only handle images if image_list is not already populated
valid_images = await self.check_images_concurrently(ptp_imagelist)
Expand All @@ -230,32 +226,28 @@ async def update_metadata_from_tracker(self, tracker_name, tracker_instance, met
else:
console.print("[yellow]Skipping PTP as no match found[/yellow]")
found_match = False
meta['skip_gen_desc'] = True
meta['description'] = None

else:
ptp_torrent_id = meta['ptp']
console.print(f"[cyan]PTP ID found in meta: {ptp_torrent_id}, using it to get IMDb ID[/cyan]")
console.print("[cyan]Using specified PTP ID to get IMDb ID[/cyan]")
imdb_id, _, meta['ext_torrenthash'] = await tracker_instance.get_imdb_from_torrent_id(ptp_torrent_id)
if imdb_id:
meta['imdb'] = str(imdb_id).zfill(7)
console.print(f"[green]IMDb ID found: tt{meta['imdb']}[/green]")
found_match = True
meta['skipit'] = True
ptp_desc, ptp_imagelist = await tracker_instance.get_ptp_description(meta['ptp'], meta, meta.get('is_disc', False))
meta['description'] = ptp_desc
meta['skip_gen_desc'] = True
if not meta.get('image_list'): # Only handle images if image_list is not already populated
valid_images = await self.check_images_concurrently(ptp_imagelist)
if valid_images:
meta['image_list'] = valid_images
console.print("[green]PTP images added to metadata.[/green]")
else:
console.print(f"[yellow]Could not find IMDb ID using PTP ID: {ptp_torrent_id}[/yellow]")
found_match = False

# Retrieve PTP description and image list
ptp_desc, ptp_imagelist = await tracker_instance.get_ptp_description(meta['ptp'], meta, meta.get('is_disc', False))
meta['description'] = ptp_desc

if not meta.get('image_list'): # Only handle images if image_list is not already populated
valid_images = await self.check_images_concurrently(ptp_imagelist)
if valid_images:
meta['image_list'] = valid_images

meta['skip_gen_desc'] = True
console.print("[green]PTP images added to metadata.[/green]")

elif tracker_name == "HDB":
if meta.get('hdb') is not None:
meta[manual_key] = meta[tracker_key]
Expand Down Expand Up @@ -3144,104 +3136,71 @@ def clean_filename(self, name):
return name

async def gen_desc(self, meta):
def clean_text(text):
return text.replace('\r\n', '').replace('\n', '').strip()

desclink = meta.get('desclink')
descfile = meta.get('descfile')
ptp_desc = ""
imagelist = []
if not meta.get('skip_gen_desc', False):
def clean_text(text):
return text.replace('\r\n', '').replace('\n', '').strip()

desc_sources = ['ptp', 'blu', 'aither', 'lst', 'oe', 'tik']
desc_source = [source.upper() for source in desc_sources if meta.get(source)]
desc_source = desc_source[0] if len(desc_source) == 1 else None

with open(f"{meta['base_dir']}/tmp/{meta['uuid']}/DESCRIPTION.txt", 'w', newline="", encoding='utf8') as description:
description.seek(0)

if (desclink, descfile, meta['desc']) == (None, None, None):
if meta.get('ptp') and str(self.config['TRACKERS'].get('PTP', {}).get('useAPI')).lower() == "true" and desc_source in ['PTP', None]:
if meta.get('skip_gen_desc', False):
console.print("[cyan]Something went wrong with PTP description.")
return meta
desclink = meta.get('desclink')
descfile = meta.get('descfile')

ptp = PTP(config=self.config)
ptp_desc, imagelist = await ptp.get_ptp_description(meta['ptp'], meta['is_disc'])
if clean_text(ptp_desc):
description.write(ptp_desc + "\n")
meta['description'] = 'PTP'
meta['imagelist'] = imagelist
with open(f"{meta['base_dir']}/tmp/{meta['uuid']}/DESCRIPTION.txt", 'w', newline="", encoding='utf8') as description:
description.seek(0)

if not ptp_desc and clean_text(meta.get('blu_desc', '')) and desc_source in ['BLU', None]:
description.write(meta['blu_desc'] + "\n")
meta['description'] = 'BLU'

if not ptp_desc and clean_text(meta.get('lst_desc', '')) and desc_source in ['LST', None]:
description.write(meta['lst_desc'] + "\n")
meta['description'] = 'LST'

if not ptp_desc and clean_text(meta.get('aither_desc', '')) and desc_source in ['AITHER', None]:
description.write(meta['aither_desc'] + "\n")
meta['description'] = 'AITHER'

if not ptp_desc and clean_text(meta.get('oe_desc', '')) and desc_source in ['OE', None]:
description.write(meta['oe_desc'] + "\n")
meta['description'] = 'OE'

if not ptp_desc and clean_text(meta.get('tike_desc', '')) and desc_source in ['TIK', None]:
description.write(meta['tik_desc'] + "\n")
meta['description'] = 'TIK'

if meta.get('desc_template'):
from jinja2 import Template
try:
with open(f"{meta['base_dir']}/data/templates/{meta['desc_template']}.txt", 'r') as f:
template = Template(f.read())
template_desc = template.render(meta)
if clean_text(template_desc):
description.write(template_desc + "\n")
console.print(f"[INFO] Description from template '{meta['desc_template']}' used.")
except FileNotFoundError:
console.print(f"[ERROR] Template '{meta['desc_template']}' not found.")

if meta.get('nfo'):
nfo_files = glob.glob("*.nfo")
if nfo_files:
nfo = nfo_files[0]
with open(nfo, 'r', encoding="utf-8") as nfo_file:
nfo_content = nfo_file.read()
description.write(f"[code]{nfo_content}[/code]\n")
if meta.get('desc_template'):
from jinja2 import Template
try:
with open(f"{meta['base_dir']}/data/templates/{meta['desc_template']}.txt", 'r') as f:
template = Template(f.read())
template_desc = template.render(meta)
if clean_text(template_desc):
description.write(template_desc + "\n")
console.print(f"[INFO] Description from template '{meta['desc_template']}' used.")
except FileNotFoundError:
console.print(f"[ERROR] Template '{meta['desc_template']}' not found.")

if meta.get('nfo'):
nfo_files = glob.glob("*.nfo")
if nfo_files:
nfo = nfo_files[0]
with open(nfo, 'r', encoding="utf-8") as nfo_file:
nfo_content = nfo_file.read()
description.write(f"[code]{nfo_content}[/code]\n")
meta['description'] = "CUSTOM"
console.print(f"[INFO] NFO file '{nfo}' used.")

if desclink:
try:
parsed = urllib.parse.urlparse(desclink.replace('/raw/', '/'))
split = os.path.split(parsed.path)
raw = parsed._replace(path=f"{split[0]}/raw/{split[1]}" if split[0] != '/' else f"/raw{parsed.path}")
raw_url = urllib.parse.urlunparse(raw)
desclink_content = requests.get(raw_url).text
description.write(desclink_content + "\n")
meta['description'] = "CUSTOM"
console.print(f"[INFO] Description from link '{desclink}' used.")
except Exception as e:
console.print(f"[ERROR] Failed to fetch description from link: {e}")

if descfile and os.path.isfile(descfile):
with open(descfile, 'r') as f:
file_content = f.read()
description.write(file_content)
meta['description'] = "CUSTOM"
console.print(f"[INFO] NFO file '{nfo}' used.")
console.print(f"[INFO] Description from file '{descfile}' used.")

if desclink:
try:
parsed = urllib.parse.urlparse(desclink.replace('/raw/', '/'))
split = os.path.split(parsed.path)
raw = parsed._replace(path=f"{split[0]}/raw/{split[1]}" if split[0] != '/' else f"/raw{parsed.path}")
raw_url = urllib.parse.urlunparse(raw)
desclink_content = requests.get(raw_url).text
description.write(desclink_content + "\n")
if meta.get('desc'):
description.write(meta['desc'] + "\n")
meta['description'] = "CUSTOM"
console.print(f"[INFO] Description from link '{desclink}' used.")
except Exception as e:
console.print(f"[ERROR] Failed to fetch description from link: {e}")

if descfile and os.path.isfile(descfile):
with open(descfile, 'r') as f:
file_content = f.read()
description.write(file_content)
meta['description'] = "CUSTOM"
console.print(f"[INFO] Description from file '{descfile}' used.")

if meta.get('desc'):
description.write(meta['desc'] + "\n")
meta['description'] = "CUSTOM"
console.print("[INFO] Custom description used.")
console.print("[INFO] Custom description used.")

description.write("\n")
description.write("\n")
return meta
else:
description_text = meta.get('description') if meta.get('description') else ""
with open(f"{meta['base_dir']}/tmp/{meta['uuid']}/DESCRIPTION.txt", 'w', newline="", encoding='utf8') as description:
description.write(description_text + "\n")

return meta
return meta

async def tag_override(self, meta):
with open(f"{meta['base_dir']}/data/tags.json", 'r', encoding="utf-8") as f:
Expand Down
7 changes: 7 additions & 0 deletions src/trackers/COMMON.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -273,12 +273,19 @@ async def unit3d_torrent_info(self, tracker, torrent_url, search_url, meta, id=N
edited_description = click.edit(description)
if edited_description:
description = edited_description.strip()
meta['description'] = description
meta['skip_gen_desc'] = True
console.print(f"[green]Final description after editing:[/green] {description}", markup=False)
elif edit_choice.lower() == 'd':
description = None
console.print("[yellow]Description discarded.[/yellow]")
else:
console.print("[green]Keeping the original description.[/green]")
meta['description'] = description
meta['skip_gen_desc'] = True
else:
meta['description'] = description
meta['skip_gen_desc'] = True

return tmdb, imdb, tvdb, mal, description, category, infohash, imagelist, file_name

Expand Down
6 changes: 5 additions & 1 deletion src/trackers/PTP.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -209,7 +209,7 @@ async def get_ptp_description(self, ptp_torrent_id, meta, is_disc):
console.print("[bold green]Successfully grabbed description from PTP")
console.print(f"[cyan]Description after cleaning:[yellow]\n{desc[:1000]}...") # Show first 1000 characters for brevity

if not meta.get('ptp') or meta['unattended']:
if not meta.get('skipit') or meta['unattended']:
# Allow user to edit or discard the description
console.print("[cyan]Do you want to edit, discard or keep the description?[/cyan]")
edit_choice = input("Enter 'e' to edit, 'd' to discard, or press Enter to keep it as is: ")
Expand All @@ -218,12 +218,16 @@ async def get_ptp_description(self, ptp_torrent_id, meta, is_disc):
edited_description = click.edit(desc)
if edited_description:
desc = edited_description.strip()
meta['description'] = desc
meta['skip_gen_desc'] = True
console.print(f"[green]Final description after editing:[/green] {desc}")
elif edit_choice.lower() == 'd':
desc = None
console.print("[yellow]Description discarded.[/yellow]")
else:
console.print("[green]Keeping the original description.[/green]")
meta['description'] = ptp_desc
meta['skip_gen_desc'] = True

return desc, imagelist

Expand Down

0 comments on commit 04f4074

Please sign in to comment.