Repository: tefra/pytuber

Issue:
Add youtube quota calculator

Discussion:

Patch:
diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml
index 8e86fd2..6164f6e 100644
--- a/.pre-commit-config.yaml
+++ b/.pre-commit-config.yaml
@@ -37,6 +37,6 @@ repos:
         args: ['--in-place', '--pre-summary-newline']
 
   - repo: https://github.com/pre-commit/mirrors-mypy
-    rev: v0.660
+    rev: v0.670
     hooks:
     -   id: mypy
diff --git a/README.rst b/README.rst
index 9f87a48..498c85a 100644
--- a/README.rst
+++ b/README.rst
@@ -56,8 +56,13 @@ Start creating youtube playlists ✨✨
 Youtube API Quota
 ~~~~~~~~~~~~~~~~~
 
-Youtube api has a daily api `quota <https://developers.google.com/youtube/v3/getting-started#quota>`_ limit.
-During development I haven't managed to reach it but I do plan for pytuber to calculate it on the fly.
+Youtube api has a daily api `quota <https://developers.google.com/youtube/v3/getting-started#quota>`_ limit which resets at midnight Pacific Time (PT).
+
+pytuber includes a quota calculator
+
+.. code-block:: console
+
+    $ pytuber quota
 
 Additionally to the api quota limit Youtube limits the amount of how many playlists you can create per day to only **10**.
 
diff --git a/pytuber/cli.py b/pytuber/cli.py
index 025dd0a..f432978 100644
--- a/pytuber/cli.py
+++ b/pytuber/cli.py
@@ -4,6 +4,7 @@ import click
 import click_completion
 
 from pytuber.core import commands as core
+from pytuber.core.utils import init_registry
 from pytuber.lastfm import commands as lastfm
 from pytuber.storage import Registry
 from pytuber.version import version
@@ -40,9 +41,7 @@ Push your updates to youtube
 @click.pass_context
 def cli(ctx: click.Context):
     cfg = os.path.join(click.get_app_dir("pytuber", False), "storage.db")
-
-    Registry.from_file(cfg)
-    Registry.set("version", version)
+    init_registry(cfg, version)
     ctx.call_on_close(lambda: Registry.persist(cfg))
     click.secho("")
 
@@ -51,6 +50,7 @@ cli.add_command(core.list)
 cli.add_command(core.show)
 cli.add_command(core.remove)
 cli.add_command(core.clean)
+cli.add_command(core.quota)
 
 
 @cli.group()
diff --git a/pytuber/core/commands/__init__.py b/pytuber/core/commands/__init__.py
index d0d6c21..27fbf7e 100644
--- a/pytuber/core/commands/__init__.py
+++ b/pytuber/core/commands/__init__.py
@@ -6,6 +6,7 @@ from pytuber.core.commands.cmd_setup import setup
 from pytuber.core.commands.cmd_show import show
 from pytuber.core.commands.cmd_autocomplete import autocomplete
 from pytuber.core.commands.cmd_clean import clean
+from pytuber.core.commands.cmd_quota import quota
 
 __all__ = [
     "setup",
@@ -16,4 +17,5 @@ __all__ = [
     "show",
     "autocomplete",
     "clean",
+    "quota",
 ]
diff --git a/pytuber/core/commands/cmd_quota.py b/pytuber/core/commands/cmd_quota.py
new file mode 100644
index 0000000..babf243
--- /dev/null
+++ b/pytuber/core/commands/cmd_quota.py
@@ -0,0 +1,35 @@
+from datetime import timedelta
+
+import click
+from tabulate import tabulate
+
+from pytuber.core.models import ConfigManager, Provider
+from pytuber.core.services import YouService
+from pytuber.utils import magenta
+
+
+@click.command()
+def quota():
+    """Show current youtube calculated quota usage."""
+
+    limit = ConfigManager.get(Provider.youtube).data["quota_limit"]
+    usage = YouService.get_quota_usage()
+    pt_date = YouService.quota_date(obj=True)
+    next_reset = timedelta(
+        hours=23 - pt_date.hour,
+        minutes=59 - pt_date.minute,
+        seconds=60 - pt_date.second,
+    )
+
+    values = [
+        (magenta("Provider:"), Provider.youtube),
+        (magenta("Limit:"), limit),
+        (magenta("Usage:"), usage),
+        (magenta("Next reset:"), str(next_reset)),
+    ]
+
+    click.secho(
+        tabulate(  # type: ignore
+            values, tablefmt="plain", colalign=("right", "left")
+        )
+    )
diff --git a/pytuber/core/commands/cmd_setup.py b/pytuber/core/commands/cmd_setup.py
index 263962f..6ee97d5 100644
--- a/pytuber/core/commands/cmd_setup.py
+++ b/pytuber/core/commands/cmd_setup.py
@@ -6,7 +6,15 @@ from pytuber.core.services import YouService
 
 @click.command("youtube")
 @click.argument("client-secrets", type=click.Path(), required=True)
-def setup(client_secrets: str) -> None:
+@click.option(
+    "--quota-limit",
+    "--q",
+    type=click.INT,
+    required=False,
+    default=1000000,
+    help="Override default youtube quota limit",
+)
+def setup(client_secrets: str, quota_limit: int) -> None:
     """
     Configure your youtube api credentials.
 
@@ -29,6 +37,7 @@ def setup(client_secrets: str) -> None:
                 client_id=credentials.client_id,
                 client_secret=credentials.client_secret,
                 scopes=credentials.scopes,
+                quota_limit=quota_limit,
             ),
         )
     )
diff --git a/pytuber/core/services.py b/pytuber/core/services.py
index b815f03..7ed474f 100644
--- a/pytuber/core/services.py
+++ b/pytuber/core/services.py
@@ -1,3 +1,5 @@
+from datetime import datetime, timedelta
+
 from google.oauth2.credentials import Credentials
 from google_auth_oauthlib.flow import InstalledAppFlow
 from googleapiclient.discovery import build
@@ -9,12 +11,14 @@ from pytuber.core.models import (
     Provider,
     Track,
 )
+from pytuber.storage import Registry
 
 
 class YouService:
     max_results = 50
     client = None
     scopes = ["https://www.googleapis.com/auth/youtube"]
+    quota_key = "youtube_quota"
 
     @classmethod
     def authorize(cls, client_secrets):
@@ -32,6 +36,7 @@ class YouService:
         )
 
         response = cls.get_client().search().list(**params).execute()
+        cls.update_quota(100)
         for item in response.get("items", []):
             if item["id"]["kind"] == "youtube#video":
                 return item["id"]["videoId"]
@@ -46,6 +51,7 @@ class YouService:
                 params.update(dict(pageToken=next_page_token))
 
             response = cls.get_client().playlists().list(**params).execute()
+            cls.update_quota(3)
             for item in response.get("items", []):
                 playlist = Playlist.from_mime(
                     item["snippet"]["description"].strip().split("\n")[-1]
@@ -71,7 +77,9 @@ class YouService:
             ),
             part="snippet,status",
         )
-        return cls.get_client().playlists().insert(**params).execute()["id"]
+        id = cls.get_client().playlists().insert(**params).execute()["id"]
+        cls.update_quota(55)
+        return id
 
     @classmethod
     def get_playlist_items(cls, playlist: Playlist):
@@ -87,6 +95,7 @@ class YouService:
                 params.update(dict(pageToken=next_page_token))
 
             resp = cls.get_client().playlistItems().list(**params).execute()
+            cls.update_quota(5)
             for item in resp.get("items", []):
 
                 try:
@@ -120,12 +129,16 @@ class YouService:
             ),
             part="snippet",
         )
-        return cls.get_client().playlistItems().insert(**params).execute()
+        result = cls.get_client().playlistItems().insert(**params).execute()
+        cls.update_quota(53)
+        return result
 
     @classmethod
     def remove_playlist_item(cls, playlist_item: PlaylistItem):
         params = dict(id=playlist_item.id)
-        return cls.get_client().playlistItems().delete(**params).execute()
+        result = cls.get_client().playlistItems().delete(**params).execute()
+        cls.update_quota(51)
+        return result
 
     @classmethod
     def get_client(cls):
@@ -136,3 +149,30 @@ class YouService:
             )
             cls.client = build("youtube", "v3", credentials=credentials)
         return cls.client
+
+    @classmethod
+    def get_quota_usage(cls):
+        return Registry.get(cls.quota_key, cls.quota_date(), default=0)
+
+    @classmethod
+    def update_quota(cls, cost: int):
+        """
+        Update current date youtube quota usage  according to this guide
+        https://developers.google.com/youtube/v3/determine_quota_cost.
+
+        :param int cost:
+        """
+        date = cls.quota_date()
+        quota = Registry.get(cls.quota_key, date, default=0) + cost
+        Registry.set(cls.quota_key, {date: quota})
+
+    @classmethod
+    def quota_date(cls, obj: bool = False):
+        """
+        Youtube daily quotas reset at midnight Pacific Time (PT). Return the
+        current quota date string.
+
+        :return: str
+        """
+        dt = datetime.utcnow() - timedelta(hours=8)
+        return dt if obj else dt.strftime("%Y%m%d")
diff --git a/pytuber/core/utils.py b/pytuber/core/utils.py
new file mode 100644
index 0000000..204fafd
--- /dev/null
+++ b/pytuber/core/utils.py
@@ -0,0 +1,14 @@
+from pytuber.storage import Registry
+
+
+def init_registry(path: str, version: str):
+    Registry.from_file(path)
+
+    current_version = Registry.get("version", default="0")
+    if current_version == "0":
+        if Registry.exists("configuration", "youtube", "data"):
+            Registry.set(
+                "configuration", "youtube", "data", "quota_limit", 1000000
+            )
+
+    Registry.set("version", version)
diff --git a/pytuber/version.py b/pytuber/version.py
index 9b80d2b..89caaf5 100644
--- a/pytuber/version.py
+++ b/pytuber/version.py
@@ -1,1 +1,1 @@
-version = "19.1"
+version = "19.2"
diff --git a/setup.py b/setup.py
index 4df6a64..8b48ece 100644
--- a/setup.py
+++ b/setup.py
@@ -18,8 +18,8 @@ if __name__ == "__main__":
             "pydrag == 18.1",
             "attrs == 18.2.0",
             "tabulate == 0.8.3",
-            "yaspin == 0.14.0",
-            "google-api-python-client == 1.7.7",
+            "yaspin == 0.14.1",
+            "google-api-python-client == 1.7.8",
             "google-auth == 1.6.2",
             "google-auth-oauthlib == 0.2.0",
         ],


Test patch:
diff --git a/tests/core/commands/test_cmd_quota.py b/tests/core/commands/test_cmd_quota.py
new file mode 100644
index 0000000..6c8784e
--- /dev/null
+++ b/tests/core/commands/test_cmd_quota.py
@@ -0,0 +1,27 @@
+from datetime import datetime
+from unittest import mock
+
+from pytuber import cli
+from pytuber.core.services import YouService
+from tests.utils import CommandTestCase, ConfigFixture
+
+
+class CommandQuotaTests(CommandTestCase):
+    @mock.patch.object(YouService, "get_quota_usage")
+    @mock.patch.object(YouService, "quota_date")
+    def test_run(self, quota_date, get_quota_usage):
+        ConfigFixture.youtube()
+        get_quota_usage.return_value = 9988
+        quota_date.return_value = datetime(
+            year=1970, month=1, day=1, hour=22, minute=22, second=11
+        )
+        result = self.runner.invoke(cli, ["quota"])
+
+        expected_output = (
+            "Provider:  youtube",
+            "     Limit:  1000000",
+            "     Usage:  9988",
+            "Next reset:  1:37:49",
+        )
+        self.assertEqual(0, result.exit_code)
+        self.assertOutput(expected_output, result.output)
diff --git a/tests/core/commands/test_cmd_setup.py b/tests/core/commands/test_cmd_setup.py
index 6c6619b..b6c745f 100644
--- a/tests/core/commands/test_cmd_setup.py
+++ b/tests/core/commands/test_cmd_setup.py
@@ -29,6 +29,7 @@ class CommandSetupYoutubeTests(CommandTestCase):
         expected = {
             "client_id": "client_id",
             "client_secret": "client_secret",
+            "quota_limit": 1000000,
             "refresh_token": None,
             "scopes": "scopes",
             "token_uri": "token_uri",
@@ -37,10 +38,12 @@ class CommandSetupYoutubeTests(CommandTestCase):
         self.assertDictEqual(expected, actual.data)
 
     def test_update(self):
-        ConfigManager.set(dict(provider=Provider.youtube, data="foo"))
+        ConfigManager.set(dict(provider=Provider.youtube, data={"foo": "bar"}))
         client_secrets = "~/Downloads/client_secrets.json"
         result = self.runner.invoke(
-            cli, ["setup", "youtube", client_secrets], input="n"
+            cli,
+            ["setup", "youtube", client_secrets, "--quota-limit", 500],
+            input="n",
         )
 
         expected_output = "\n".join(
diff --git a/tests/core/test_services.py b/tests/core/test_services.py
index 7f6bfac..df47e8c 100644
--- a/tests/core/test_services.py
+++ b/tests/core/test_services.py
@@ -1,3 +1,4 @@
+from datetime import datetime, timedelta
 from unittest import mock
 
 from google.oauth2.credentials import Credentials
@@ -46,6 +47,7 @@ class YouServiceTests(TestCase):
             q="{} {}".format(track.artist, track.name),
             type="video",
         )
+        self.assertEqual(100, YouService.get_quota_usage())
 
     @mock.patch.object(YouService, "get_client")
     def test_get_playlists(self, get_client):
@@ -81,6 +83,7 @@ class YouServiceTests(TestCase):
                 mock.call().execute(),
             ]
         )
+        self.assertEqual(6, YouService.get_quota_usage())
 
     @mock.patch.object(YouService, "get_client")
     def test_create_playlist(self, get_client):
@@ -96,6 +99,7 @@ class YouServiceTests(TestCase):
             ),
             part="snippet,status",
         )
+        self.assertEqual(55, YouService.get_quota_usage())
 
     @mock.patch.object(YouService, "get_client")
     def test_get_playlist_items(self, get_client):
@@ -153,6 +157,7 @@ class YouServiceTests(TestCase):
                 mock.call().execute(),
             ]
         )
+        self.assertEqual(10, YouService.get_quota_usage())
 
     @mock.patch.object(YouService, "get_client")
     def test_create_playlist_item(self, get_client):
@@ -172,6 +177,7 @@ class YouServiceTests(TestCase):
             ),
             part="snippet",
         )
+        self.assertEqual(53, YouService.get_quota_usage())
 
     @mock.patch.object(YouService, "get_client")
     def test_remove_playlist_item(self, get_client):
@@ -181,6 +187,7 @@ class YouServiceTests(TestCase):
 
         self.assertEqual("foo", YouService.remove_playlist_item(item))
         delete.assert_called_once_with(id=item.id)
+        self.assertEqual(51, YouService.get_quota_usage())
 
     @mock.patch("pytuber.core.services.build")
     @mock.patch.object(Credentials, "from_authorized_user_info")
@@ -197,3 +204,7 @@ class YouServiceTests(TestCase):
 
         get_user_info.assert_called_once_with("foo", scopes=YouService.scopes)
         build.assert_called_once_with("youtube", "v3", credentials="creds")
+
+    def test_quota_date(self):
+        expected = (datetime.utcnow() - timedelta(hours=8)).strftime("%Y%m%d")
+        self.assertEqual(expected, YouService.quota_date())
diff --git a/tests/utils.py b/tests/utils.py
index 683d6a8..1abf153 100644
--- a/tests/utils.py
+++ b/tests/utils.py
@@ -5,7 +5,13 @@ from unittest import mock
 
 from click.testing import CliRunner
 
-from pytuber.core.models import Playlist, PlaylistItem, Track
+from pytuber.core.models import (
+    ConfigManager,
+    Playlist,
+    PlaylistItem,
+    Provider,
+    Track,
+)
 from pytuber.storage import Registry
 
 
@@ -103,3 +109,21 @@ class PlaylistItemFixture(Fixture):
         )
         params.update(kwargs)
         return params
+
+
+class ConfigFixture:
+    @classmethod
+    def youtube(cls):
+        ConfigManager.set(
+            dict(
+                provider=Provider.youtube.value,
+                data=dict(
+                    refresh_token=None,
+                    token_uri=None,
+                    client_id=None,
+                    client_secret=None,
+                    scopes=None,
+                    quota_limit=100,
+                ),
+            )
+        )