diff --git a/src/file_classifier.py b/src/file_classifier.py index 37c1b39..a54b780 100644 --- a/src/file_classifier.py +++ b/src/file_classifier.py @@ -26,6 +26,43 @@ def get_metadata(filepath: Path) -> dict: return json.loads(result.stdout) +def classify_show(info: PathInfo) -> PathInfo: + # Prepare meta information + episode_durations: list[int] = [] + movie_information = None + if info.category == PathCategory.SHOW: + for ep in imdb.get_all_episodes(info.imdb.imdb_id): + if ep.duration: + episode_durations.append(int(ep.duration / 60)) + + print(set(episode_durations)) + # Go over all files + count = 0 + for root, dirs, files in os.walk(info.path): + for filename in files: + filepath = Path(os.path.join(root, filename)) + if filepath.suffix != ".mkv": + continue + m = get_metadata(filepath) + duration = int(float(m["format"]["duration"]) / 60) + if ( + math.floor(duration) in episode_durations + or math.ceil(duration) in episode_durations + ): + print(f"{filename} {duration} EPISODE") + count += 1 + else: + print(f"{filename} {duration}") + + print(count) + return info + + +def classify_movie(info: PathInfo) -> PathInfo: + logging.error(f"Movie classification not yet implemented.") + return info + + def classify_files(path: str) -> PathInfo | None: p = Path(path) @@ -61,43 +98,25 @@ def classify_files(path: str) -> PathInfo | None: return None logging.info(f"Found matching IMDB entry with id [{imdb_entry.imdb_id}].") - info = PathInfo(path=p, title=title, year=year, imdb_id=imdb_entry.imdb_id) + info = PathInfo( + path=p, title=title, year=year, imdb_id=imdb_entry.imdb_id, imdb=imdb_entry + ) # Identify category if imdb_entry.kind == "tvSeries": info.category = PathCategory.SHOW + logging.info(f"Path identified as containing SHOW.") + info = classify_show(info) elif imdb_entry.kind == "movie": info.category = PathCategory.MOVIE + logging.info(f"Path identified as containing MOVIE.") + info = classify_movie(info) else: info.category = PathCategory.UNKNOWN + logging.error( + f"IMDB entry has unknown qualifier for content [{imdb_entry.kind}]." + ) - # Prepare meta information - episode_durations: list[int] = [] - movie_information = None - if info.category == PathCategory.SHOW: - for ep in imdb.get_all_episodes(imdb_entry.imdb_id): - episode_durations.append(int(ep.duration / 60)) - - print(set(episode_durations)) - # Go over all files - count = 0 - for root, dirs, files in os.walk(p): - for filename in files: - filepath = Path(os.path.join(root, filename)) - if filepath.suffix != ".mkv": - continue - m = get_metadata(filepath) - duration = int(float(m["format"]["duration"]) / 60) - if ( - math.floor(duration) in episode_durations - or math.ceil(duration) in episode_durations - ): - print(f"{filename} {duration} EPISODE") - count += 1 - else: - print(f"{filename} {duration}") - - print(count) return info diff --git a/src/structures.py b/src/structures.py index f4a7f4f..9a62c34 100644 --- a/src/structures.py +++ b/src/structures.py @@ -1,6 +1,7 @@ from dataclasses import dataclass, field from enum import Enum from pathlib import Path +from imdbinfo.models import MovieBriefInfo from os import path @@ -89,6 +90,9 @@ class PathInfo: imdb_id: str """IMDB id.""" + imdb: MovieBriefInfo + """IMDB info object referencing media.""" + category: PathCategory = PathCategory.UNCLASSIFIED """Category of path media."""