Skip to content

Commit

Permalink
Improved BLU description handling
Browse files Browse the repository at this point in the history
  • Loading branch information
@Audionut
Audionut committed Sep 1, 2024
1 parent e42c87e commit 67ef332
Show file tree
Hide file tree
Showing 2 changed files with 119 additions and 122 deletions.
99 changes: 45 additions & 54 deletions src/prep.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -90,11 +90,7 @@ async def update_metadata_from_tracker(self, tracker_name, tracker_instance, met
manual_key = f"{tracker_key}_manual"
found_match = False

# console.print(f"[cyan]Starting update_metadata_from_tracker for: {tracker_name}[/cyan]")

# Handle each tracker separately
if tracker_name == "BLU":
# console.print(f"[blue]Handling BLU tracker[/blue]")
if meta.get(tracker_key) is not None:
console.print(f"[cyan]{tracker_name} ID found in meta, reusing existing ID: {meta[tracker_key]}[/cyan]")
blu_tmdb, blu_imdb, blu_tvdb, blu_mal, blu_desc, blu_category, meta['ext_torrenthash'], blu_imagelist, blu_filename = await COMMON(self.config).unit3d_torrent_info(
Expand All @@ -103,33 +99,29 @@ async def update_metadata_from_tracker(self, tracker_name, tracker_instance, met
tracker_instance.search_url,
id=meta[tracker_key]
)
# console.print(f"[blue]BLU search by ID complete[/blue]")
if blu_tmdb not in [None, '0'] or blu_imdb not in [None, '0'] or blu_tvdb not in [None, '0']:
console.print(f"[green]Valid data found on {tracker_name}, setting meta values[/green]")
if await self.prompt_user_for_id_selection(blu_tmdb, blu_imdb, blu_tvdb, blu_filename):
if blu_tmdb not in [None, '0']:
meta['tmdb_manual'] = blu_tmdb
if blu_imdb not in [None, '0']:
meta['imdb'] = str(blu_imdb).zfill(7) # Pad IMDb ID with leading zeros
if blu_tvdb not in [None, '0']:
meta['tvdb_id'] = blu_tvdb
if blu_mal not in [None, '0']:
meta['mal'] = blu_mal
if blu_desc not in [None, '0', '']:
meta['blu_desc'] = blu_desc
if blu_category.upper() in ['MOVIE', 'TV SHOW', 'FANRES']:
meta['category'] = 'TV' if blu_category.upper() == 'TV SHOW' else blu_category.upper()
if not meta.get('image_list'): # Only handle images if image_list is not already populated
if blu_imagelist: # Ensure blu_imagelist is not empty before setting
meta['image_list'] = blu_imagelist
if meta.get('image_list'): # Double-check if image_list is set before handling it
await self.handle_image_list(meta, tracker_name)
if blu_filename:
meta['blu_filename'] = blu_filename # Store the filename in meta for later use
found_match = True
console.print("[green]BLU data successfully updated in meta[/green]")
else:
console.print(f"[yellow]Skipped {tracker_name}, moving to the next site.[/yellow]")
if blu_tmdb not in [None, '0']:
meta['tmdb_manual'] = blu_tmdb
if blu_imdb not in [None, '0']:
meta['imdb'] = str(blu_imdb).zfill(7) # Pad IMDb ID with leading zeros
if blu_tvdb not in [None, '0']:
meta['tvdb_id'] = blu_tvdb
if blu_mal not in [None, '0']:
meta['mal'] = blu_mal
if blu_desc not in [None, '0', '']:
meta['blu_desc'] = blu_desc
if blu_category.upper() in ['MOVIE', 'TV SHOW', 'FANRES']:
meta['category'] = 'TV' if blu_category.upper() == 'TV SHOW' else blu_category.upper()
if not meta.get('image_list'): # Only handle images if image_list is not already populated
if blu_imagelist: # Ensure blu_imagelist is not empty before setting
meta['image_list'] = blu_imagelist
if meta.get('image_list'): # Double-check if image_list is set before handling it
await self.handle_image_list(meta, tracker_name)
if blu_filename:
meta['blu_filename'] = blu_filename # Store the filename in meta for later use
found_match = True
console.print("[green]BLU data successfully updated in meta[/green]")
else:
console.print(f"[yellow]No valid data found on {tracker_name}[/yellow]")
else:
Expand All @@ -140,33 +132,32 @@ async def update_metadata_from_tracker(self, tracker_name, tracker_instance, met
tracker_instance.search_url,
file_name=search_term
)
# console.print(f"[blue]BLU search by file name complete[/blue]")

if blu_tmdb not in [None, '0'] or blu_imdb not in [None, '0'] or blu_tvdb not in [None, '0']:
console.print(f"[green]Valid data found on {tracker_name} using file name, setting meta values[/green]")
if await self.prompt_user_for_id_selection(blu_tmdb, blu_imdb, blu_tvdb, blu_filename):
if blu_tmdb not in [None, '0']:
meta['tmdb_manual'] = blu_tmdb
if blu_imdb not in [None, '0']:
meta['imdb'] = str(blu_imdb).zfill(7)
if blu_tvdb not in [None, '0']:
meta['tvdb_id'] = blu_tvdb
if blu_mal not in [None, '0']:
meta['mal'] = blu_mal
if blu_desc not in [None, '0', '']:
meta['blu_desc'] = blu_desc
if blu_category.upper() in ['MOVIE', 'TV SHOW', 'FANRES']:
meta['category'] = 'TV' if blu_category.upper() == 'TV SHOW' else blu_category.upper()
if not meta.get('image_list'): # Only handle images if image_list is not already populated
if blu_imagelist: # Ensure blu_imagelist is not empty before setting
meta['image_list'] = blu_imagelist
if meta.get('image_list'): # Double-check if image_list is set before handling it
await self.handle_image_list(meta, tracker_name)
if blu_filename:
meta['blu_filename'] = blu_filename
found_match = True
console.print("[green]BLU data successfully updated in meta[/green]")
else:
console.print(f"[yellow]Skipped {tracker_name}, moving to the next site.[/yellow]")

if blu_tmdb not in [None, '0']:
meta['tmdb_manual'] = blu_tmdb
if blu_imdb not in [None, '0']:
meta['imdb'] = str(blu_imdb).zfill(7)
if blu_tvdb not in [None, '0']:
meta['tvdb_id'] = blu_tvdb
if blu_mal not in [None, '0']:
meta['mal'] = blu_mal
if blu_desc not in [None, '0', '']:
meta['blu_desc'] = blu_desc
if blu_category.upper() in ['MOVIE', 'TV SHOW', 'FANRES']:
meta['category'] = 'TV' if blu_category.upper() == 'TV SHOW' else blu_category.upper()
if not meta.get('image_list'): # Only handle images if image_list is not already populated
if blu_imagelist: # Ensure blu_imagelist is not empty before setting
meta['image_list'] = blu_imagelist
if meta.get('image_list'): # Double-check if image_list is set before handling it
await self.handle_image_list(meta, tracker_name)
if blu_filename:
meta['blu_filename'] = blu_filename

found_match = True
console.print("[green]BLU data successfully updated in meta[/green]")
else:
console.print(f"[yellow]No valid data found on {tracker_name}[/yellow]")

Expand Down
142 changes: 74 additions & 68 deletions src/trackers/COMMON.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -142,6 +142,26 @@ async def unit3d_distributor_ids(self, distributor):
}.get(distributor, 0)
return distributor_id

async def prompt_user_for_id_selection(self, blu_tmdb=None, blu_imdb=None, blu_tvdb=None, blu_filename=None, imdb=None):
if imdb:
imdb = str(imdb).zfill(7) # Convert to string and ensure IMDb ID is 7 characters long by adding leading zeros
console.print(f"[cyan]Found IMDb ID: https://www.imdb.com/title/tt{imdb}[/cyan]")
if blu_tmdb or blu_imdb or blu_tvdb:
if blu_imdb:
blu_imdb = str(blu_imdb).zfill(7) # Convert to string and ensure IMDb ID is 7 characters long by adding leading zeros
console.print("[cyan]Found the following IDs on BLU:")
console.print(f"TMDb ID: {blu_tmdb}")
console.print(f"IMDb ID: https://www.imdb.com/title/tt{blu_imdb}")
console.print(f"TVDb ID: {blu_tvdb}")
console.print(f"Filename: {blu_filename}") # Ensure filename is printed if available

selection = input("Do you want to use this ID? (y/n): ").strip().lower()
return selection == 'y'

async def prompt_user_for_confirmation(self, message):
selection = input(f"{message} (y/n): ").strip().lower()
return selection == 'y'

async def unit3d_torrent_info(self, tracker, torrent_url, search_url, id=None, file_name=None):
tmdb = imdb = tvdb = description = category = infohash = mal = files = None # noqa F841
imagelist = []
Expand All @@ -161,14 +181,13 @@ async def unit3d_torrent_info(self, tracker, torrent_url, search_url, id=None, f
return None, None, None, None, None, None, None, None, None

response = requests.get(url=url, params=params)
# console.print(f"Requested URL: {response.url}")
# console.print(f"Status Code: {response.status_code}")

try:
json_response = response.json()
# console.print(json_response)

# console.print(f"[blue]Raw API Response: {json_response}[/blue]")

except ValueError:
# console.print(f"Response Text: {response.text}")
return None, None, None, None, None, None, None, None, None

try:
Expand All @@ -185,76 +204,63 @@ async def unit3d_torrent_info(self, tracker, torrent_url, search_url, id=None, f
mal = attributes.get('mal_id')
imdb = attributes.get('imdb_id')
infohash = attributes.get('info_hash')

if description:
bbcode = BBCODE()
description, imagelist = bbcode.clean_unit3d_description(description, torrent_url)
console.print(f"[green]Successfully grabbed description from {tracker}")
console.print(f"[blue]Extracted description: [yellow]{description}")

# 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("[cyan]Enter 'e' to edit, 'd' to discard, or press Enter to keep it as is: [/cyan]")

if edit_choice.lower() == 'e':
edited_description = click.edit(description)
if edited_description:
description = edited_description.strip()
console.print(f"[green]Final description after editing:[/green] {description}")
elif edit_choice.lower() == 'd':
description = None
console.print("[yellow]Description discarded.[/yellow]")
else:
console.print("[green]Keeping the original description.[/green]")
else:
console.print(f"[yellow]No description found for {tracker}.[/yellow]")
else:
console.print(f"[yellow]No data found in the response for {tracker} when searching by file name.[/yellow]")

# Handle response when searching by ID
if id and not data:
attributes = json_response.get('attributes', {})

# Extract data from the attributes
category = attributes.get('category')
description = attributes.get('description')
tmdb = attributes.get('tmdb_id')
tvdb = attributes.get('tvdb_id')
mal = attributes.get('mal_id')
imdb = attributes.get('imdb_id')
infohash = attributes.get('info_hash')

if description:
bbcode = BBCODE()
description, imagelist = bbcode.clean_unit3d_description(description, torrent_url)
console.print(f"[green]Successfully grabbed description from {tracker}")
console.print(f"[blue]Extracted description: [yellow]{description}")

# 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("[cyan]Enter 'e' to edit, 'd' to discard, or press Enter to keep it as is: [/cyan]")

if edit_choice.lower() == 'e':
edited_description = click.edit(description)
if edited_description:
description = edited_description.strip()
console.print(f"[green]Final description after editing:[/green] {description}")
elif edit_choice.lower() == 'd':
description = None
console.print("[yellow]Description discarded.[/yellow]")
else:
console.print("[green]Keeping the original description.[/green]")
# Handle response when searching by ID
if id and not data:
attributes = json_response.get('attributes', {})

# Extract data from the attributes
category = attributes.get('category')
description = attributes.get('description')
tmdb = attributes.get('tmdb_id')
tvdb = attributes.get('tvdb_id')
mal = attributes.get('mal_id')
imdb = attributes.get('imdb_id')
infohash = attributes.get('info_hash')

# Handle file name extraction
files = attributes.get('files', [])
if files:
if len(files) == 1:
file_name = files[0]['name']
else:
file_name = [file['name'] for file in files[:5]] # Return up to 5 filenames

console.print(f"[blue]Extracted filename(s): {file_name}[/blue]") # Print the extracted filename(s)

if tmdb or imdb or tvdb:
console.print(f"[green]Valid IDs found: TMDb: {tmdb}, IMDb: {imdb}, TVDb: {tvdb}[/green]")
if not await self.prompt_user_for_id_selection(tmdb, imdb, tvdb, file_name):
console.print("[yellow]User chose to skip based on IDs.[/yellow]")
return None, None, None, None, None, None, None, None, None

if description:
bbcode = BBCODE()
description, imagelist = bbcode.clean_unit3d_description(description, torrent_url)
console.print(f"[green]Successfully grabbed description from {tracker}")
console.print(f"[blue]Extracted description: [yellow]{description}")

# 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("[cyan]Enter 'e' to edit, 'd' to discard, or press Enter to keep it as is: [/cyan]")

if edit_choice.lower() == 'e':
edited_description = click.edit(description)
if edited_description:
description = edited_description.strip()
console.print(f"[green]Final description after editing:[/green] {description}")
elif edit_choice.lower() == 'd':
description = None
console.print("[yellow]Description discarded.[/yellow]")
else:
console.print(f"[yellow]No description found for {tracker}.[/yellow]")
console.print("[green]Keeping the original description.[/green]")

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

except Exception as e:
console.print_exception()
console.print(f"[yellow]Invalid Response from {tracker} API. Error: {str(e)}[/yellow]")

if description: # Ensure description is only printed if it's not None
console.print(f"[green]Final description to be returned:[/green] {description}")

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

async def parseCookieFile(self, cookiefile):
"""Parse a cookies.txt file and return a dictionary of key value pairs
Expand Down

0 comments on commit 67ef332

Please sign in to comment.