Add rich tree generator to sorted title wrappers

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)

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)

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)