From e28ae87bbe638cab7d1a960bdf6af520ee2eedf1 Mon Sep 17 00:00:00 2001
From: Lars Noack <lars@webcontact.de>
Date: Wed, 24 Apr 2024 12:35:22 +0200
Subject: [PATCH] feat: moved id parsing in class

---
 python_sponsorblock/__init__.py | 49 +++++++++++++++++----------------
 tests/test_video_id.py          | 23 ++++++++++------
 2 files changed, 41 insertions(+), 31 deletions(-)

diff --git a/python_sponsorblock/__init__.py b/python_sponsorblock/__init__.py
index 63f00f6..8388ece 100644
--- a/python_sponsorblock/__init__.py
+++ b/python_sponsorblock/__init__.py
@@ -2,30 +2,10 @@ import requests
 import logging
 from urllib.parse import urlparse, urlunparse, parse_qs
 import re
-from typing import Optional
+from typing import Optional, List
 
 from .exceptions import SponsorBlockError, SponsorBlockIdNotFoundError
-
-
-def _get_video_id(i: str, silent: bool = False) -> None:
-    # check with regex if i is already an id like r1Fa1iWJVEA
-    if re.match(r"^[a-zA-Z0-9_-]{11}$", i):
-        return i.strip()
-
-    url = urlparse(url=i)
-
-    if url.netloc == "youtu.be":
-        return url.path[1:]
-
-    type_frag_list = url.path.split("/")
-            
-    query_stuff = parse_qs(url.query)
-    if "v" not in query_stuff:
-        if not silent:
-            raise SponsorBlockIdNotFoundError("No video id found in the url")
-        return None
-    else:
-        return query_stuff["v"][0]
+from .constants import Segment
 
 class SponsorBlock:
     def __init__(self, session: requests.Session = None, base_url: str = "https://sponsor.ajay.app", silent: bool = False, _requests_logging_exists: bool = False):
@@ -37,6 +17,25 @@ class SponsorBlock:
 
         self.logger: logging.Logger = logging.Logger("SponsorBlock")
 
+    def _get_video_id(self, video: str) -> Optional[str]:
+        if re.match(r"^[a-zA-Z0-9_-]{11}$", video):
+            return video.strip()
+
+        url = urlparse(url=video)
+
+        if url.netloc == "youtu.be":
+            return url.path[1:]
+
+        type_frag_list = url.path.split("/")
+                
+        query_stuff = parse_qs(url.query)
+        if "v" not in query_stuff:
+            if not self.silent:
+                raise SponsorBlockIdNotFoundError("No video id found in the url")
+            return None
+        else:
+            return query_stuff["v"][0]
+
     def _request(self, method: str, endpoint: str) -> Optional[requests.Response]:
         error_message = ""
         url = self.base_url + endpoint
@@ -55,4 +54,8 @@ class SponsorBlock:
             if not self.silent:
                 raise exceptions.SponsorBlockConnectionError(error_message)
 
-        return r
\ No newline at end of file
+        return r
+
+    def get_segments(video: str) -> List[Segment]:
+        video_id = _get_video_id
+        r: List[Segment] = []
diff --git a/tests/test_video_id.py b/tests/test_video_id.py
index 9016268..d03285b 100644
--- a/tests/test_video_id.py
+++ b/tests/test_video_id.py
@@ -1,6 +1,6 @@
 import unittest
 
-from python_sponsorblock import _get_video_id, exceptions
+from python_sponsorblock import SponsorBlock, exceptions
 
 
 MAGIC_ID = "dQw4w9WgXcQ"
@@ -8,25 +8,32 @@ MAGIC_ID = "dQw4w9WgXcQ"
 
 class TestVideoIdParsing(unittest.TestCase):
     def test_already_parsed_id(self):
-        self.assertEqual(_get_video_id(MAGIC_ID), MAGIC_ID)
+        s = SponsorBlock()
+        self.assertEqual(s._get_video_id(MAGIC_ID), MAGIC_ID)
 
     def test_non_existing_id(self):
+        s = SponsorBlock()
         with self.assertRaises(exceptions.SponsorBlockIdNotFoundError):
-            _get_video_id("eeeeeeeeeeeeeeeeeeeeeeeeeeeeeeee")
+            s._get_video_id("eeeeeeeeeeeeeeeeeeeeeeeeeeeeeeee")
 
-        self.assertEqual(_get_video_id("eeeeeeeeeeeeeeeeeeeeeeeeeeeeeeee", silent=True), None)
+        s = SponsorBlock(silent=True)
+        self.assertEqual(s._get_video_id("eeeeeeeeeeeeeeeeeeeeeeeeeeeeeeee"), None)
 
     def test_youtu_dot_be(self):
-        self.assertEqual(_get_video_id("https://youtu.be/" + MAGIC_ID), MAGIC_ID)
+        s = SponsorBlock()
+        self.assertEqual(s._get_video_id("https://youtu.be/" + MAGIC_ID), MAGIC_ID)
 
     def test_youtube_dot_com(self):
-        self.assertEqual(_get_video_id("https://www.youtube.com/watch?v=" + MAGIC_ID), MAGIC_ID)
+        s = SponsorBlock()
+        self.assertEqual(s._get_video_id("https://www.youtube.com/watch?v=" + MAGIC_ID), MAGIC_ID)
 
     def test_invidious(self):
-        self.assertEqual(_get_video_id("https://invidio.us/watch?v=" + MAGIC_ID), MAGIC_ID)
+        s = SponsorBlock()
+        self.assertEqual(s._get_video_id("https://invidio.us/watch?v=" + MAGIC_ID), MAGIC_ID)
     
     def test_piped(self):
-        self.assertEqual(_get_video_id("https://piped.video/watch?v=" + MAGIC_ID), MAGIC_ID)
+        s = SponsorBlock()
+        self.assertEqual(s._get_video_id("https://piped.video/watch?v=" + MAGIC_ID), MAGIC_ID)
 
 
 if __name__ == "__main__":