From f3c06ec15c041737400f7e4aacf73db0117b6063 Mon Sep 17 00:00:00 2001
From: Jesse Bannon <jbann1994@gmail.com>
Date: Thu, 26 May 2022 23:54:12 -0700
Subject: [PATCH] Add more source variables (#42)

---
 .../entries/variables/entry_variables.py      |  9 +++
 .../entries/variables/youtube_variables.py    | 50 ++++++++++++
 tests/unit/entries/conftest.py                |  1 +
 tests/unit/entries/test_youtube_entries.py    | 79 +++++++++++++++++++
 4 files changed, 139 insertions(+)
 create mode 100644 tests/unit/entries/test_youtube_entries.py

diff --git a/src/ytdl_sub/entries/variables/entry_variables.py b/src/ytdl_sub/entries/variables/entry_variables.py
index 364aaba6..d9999d95 100644
--- a/src/ytdl_sub/entries/variables/entry_variables.py
+++ b/src/ytdl_sub/entries/variables/entry_variables.py
@@ -105,6 +105,15 @@ class EntryVariables(SourceVariables):
         """
         return int(self.upload_date[:4])
 
+    @property
+    def upload_year_truncated(self) -> int:
+        """
+        Returns
+        -------
+        The last two digits of the upload year, i.e. 22 in 2022
+        """
+        return int(self.upload_date[:2])
+
     @property
     def upload_month_padded(self) -> str:
         """
diff --git a/src/ytdl_sub/entries/variables/youtube_variables.py b/src/ytdl_sub/entries/variables/youtube_variables.py
index bdace313..1f68d7af 100644
--- a/src/ytdl_sub/entries/variables/youtube_variables.py
+++ b/src/ytdl_sub/entries/variables/youtube_variables.py
@@ -1,3 +1,5 @@
+from yt_dlp.utils import sanitize_filename
+
 from ytdl_sub.entries.base_entry import BaseEntry
 from ytdl_sub.entries.variables.entry_variables import EntryVariables
 
@@ -6,6 +8,24 @@ from ytdl_sub.entries.variables.entry_variables import EntryVariables
 
 
 class YoutubeVideoVariables(EntryVariables):
+    @property
+    def channel(self: BaseEntry) -> str:
+        """
+        Returns
+        -------
+        The channel name.
+        """
+        return self.kwargs("channel")
+
+    @property
+    def channel_sanitized(self) -> str:
+        """
+        Returns
+        -------
+        The channel name, sanitized.
+        """
+        return sanitize_filename(self.channel)
+
     @property
     def track_title(self: BaseEntry) -> str:
         """
@@ -19,6 +39,36 @@ class YoutubeVideoVariables(EntryVariables):
 
         return super().title
 
+    @property
+    def track_title_sanitized(self) -> str:
+        """
+        Returns
+        -------
+        The sanitized track title.
+        """
+        return sanitize_filename(self.track_title)
+
+    @property
+    def artist(self: BaseEntry) -> str:
+        """
+        Returns
+        -------
+        The artist of a music video if it is available, otherwise it falls back to the channel.
+        """
+        if self.kwargs_contains("artist"):
+            return self.kwargs("artist")
+
+        return self.kwargs("channel")
+
+    @property
+    def artist_sanitized(self) -> str:
+        """
+        Returns
+        -------
+        The sanitized artist name.
+        """
+        return sanitize_filename(self.artist)
+
     @property
     def playlist_index(self) -> int:
         """
diff --git a/tests/unit/entries/conftest.py b/tests/unit/entries/conftest.py
index 60217f18..e24a9aef 100644
--- a/tests/unit/entries/conftest.py
+++ b/tests/unit/entries/conftest.py
@@ -92,6 +92,7 @@ def mock_entry_to_dict(
         "upload_date": upload_date,
         "upload_date_standardized": f"{upload_year}-{_pad(upload_month)}-{_pad(upload_day)}",
         "upload_year": upload_year,
+        "upload_year_truncated": int(str(upload_year)[:2]),
         "upload_month": upload_month,
         "upload_month_padded": _pad(upload_month),
         "upload_day": upload_day,
diff --git a/tests/unit/entries/test_youtube_entries.py b/tests/unit/entries/test_youtube_entries.py
new file mode 100644
index 00000000..5335d619
--- /dev/null
+++ b/tests/unit/entries/test_youtube_entries.py
@@ -0,0 +1,79 @@
+import pytest
+
+from ytdl_sub.entries.soundcloud import SoundcloudTrack
+from ytdl_sub.entries.youtube import YoutubeVideo
+
+
+@pytest.fixture
+def playlist_index():
+    return 1
+
+
+@pytest.fixture
+def playlist_size():
+    return 1
+
+
+@pytest.fixture
+def channel():
+    return "the channel"
+
+
+@pytest.fixture
+def track_title():
+    return "not the title!"
+
+
+@pytest.fixture
+def artist():
+    return "not the channel"
+
+
+@pytest.fixture
+def mock_youtube_video_to_dict(
+    mock_entry_to_dict, playlist_index, playlist_size, channel, track_title, artist
+):
+    return dict(
+        mock_entry_to_dict,
+        **{
+            "playlist_index": playlist_index,
+            "playlist_size": playlist_size,
+            "channel": channel,
+            "channel_sanitized": channel,
+            "track_title": track_title,
+            "track_title_sanitized": track_title,
+            "artist": artist,
+            "artist_sanitized": artist,
+        }
+    )
+
+
+@pytest.fixture
+def mock_youtube_video_kwargs(
+    mock_entry_kwargs, playlist_index, playlist_size, channel, track_title, artist
+):
+    return dict(
+        mock_entry_kwargs,
+        **{
+            "playlist_index": playlist_index,
+            "playlist_count": playlist_size,
+            "channel": channel,
+            "track": track_title,
+            "artist": artist,
+        }
+    )
+
+
+@pytest.fixture
+def mock_youtube_video(mock_youtube_video_kwargs):
+    return YoutubeVideo(entry_dict=mock_youtube_video_kwargs)
+
+
+class TestYoutubeVideo(object):
+    def test_to_dict(self, mock_youtube_video, mock_youtube_video_to_dict):
+        assert mock_youtube_video.to_dict() == mock_youtube_video_to_dict
+
+    def test_youtube_dict_contains_valid_formatters(
+        self, mock_youtube_video, validate_entry_dict_contains_valid_formatters
+    ):
+        assert validate_entry_dict_contains_valid_formatters(mock_youtube_video)