diff --git a/devine/core/titles/episode.py b/devine/core/titles/episode.py index 37f36d5..4ed4a4c 100644 --- a/devine/core/titles/episode.py +++ b/devine/core/titles/episode.py @@ -1,9 +1,11 @@ import re from abc import ABC +from collections import Counter from typing import Any, Iterable, Optional, Union from langcodes import Language from pymediainfo import MediaInfo +from rich.tree import Tree from sortedcontainers import SortedKeyList from devine.core.config import config @@ -178,5 +180,29 @@ class Series(SortedKeyList, ABC): return super().__str__() return self[0].title + (f" ({self[0].year})" if self[0].year else "") + def tree(self, verbose: bool = False) -> Tree: + seasons = Counter(x.season for x in self) + num_seasons = len(seasons) + num_episodes = sum(seasons.values()) + tree = Tree( + f"{num_seasons} Season{['s', ''][num_seasons == 1]}, {num_episodes} Episode{['s', ''][num_episodes == 1]}", + guide_style="bright_black" + ) + if verbose: + for season, episodes in seasons.items(): + season_tree = tree.add( + f"[bold]Season {str(season).zfill(len(str(num_seasons)))}[/]: [bright_black]{episodes} episodes", + guide_style="bright_black" + ) + for episode in self: + if episode.season == season: + season_tree.add( + f"[bold]{str(episode.number).zfill(len(str(episodes)))}" + ( + f".[/] [bright_black]{episode.name}" + ) if episode.name else "" + ) + + return tree + __ALL__ = (Episode, Series) diff --git a/devine/core/titles/movie.py b/devine/core/titles/movie.py index f47ce24..4c4a4f3 100644 --- a/devine/core/titles/movie.py +++ b/devine/core/titles/movie.py @@ -3,6 +3,7 @@ from typing import Any, Iterable, Optional, Union from langcodes import Language from pymediainfo import MediaInfo +from rich.tree import Tree from sortedcontainers import SortedKeyList from devine.core.config import config @@ -136,5 +137,19 @@ class Movies(SortedKeyList, ABC): # TODO: Assumes there's only one movie return self[0].name + (f" ({self[0].year})" if self[0].year else "") + def tree(self, *_) -> Tree: + num_movies = len(self) + tree = Tree( + f"{num_movies} Movie{['s', ''][num_movies == 1]}", + guide_style="bright_black" + ) + for movie in self: + tree.add( + f"[bold]{movie.name}[/] [bright_black]({movie.year or '?'})", + guide_style="bright_black" + ) + + return tree + __ALL__ = (Movie, Movies) diff --git a/devine/core/titles/song.py b/devine/core/titles/song.py index 04f1652..48236b8 100644 --- a/devine/core/titles/song.py +++ b/devine/core/titles/song.py @@ -3,6 +3,7 @@ from typing import Any, Iterable, Optional, Union from langcodes import Language from pymediainfo import MediaInfo +from rich.tree import Tree from sortedcontainers import SortedKeyList from devine.core.config import config @@ -131,5 +132,19 @@ class Album(SortedKeyList, ABC): return super().__str__() return f"{self[0].artist} - {self[0].album} ({self[0].year or '?'})" + def tree(self, *_) -> Tree: + num_songs = len(self) + tree = Tree( + f"{num_songs} Song{['s', ''][num_songs == 1]}", + guide_style="bright_black" + ) + for song in self: + tree.add( + f"[bold]Track {song.track:02}.[/] [bright_black]({song.name})", + guide_style="bright_black" + ) + + return tree + __ALL__ = (Song, Album)