Merge pull request #4 from hbmartin/regex-refactor

pytube/__main__.py

@@ -175,7 +175,8 @@ class YouTube:
             or '<img class="icon meh" src="/yts/img'  # noqa: W503
             not in self.watch_html  # noqa: W503
         ):
-            raise VideoUnavailable("This video is unavailable.")
+            raise VideoUnavailable(video_id=self.video_id)
+
         self.embed_html = request.get(url=self.embed_url)
         self.age_restricted = extract.is_age_restricted(self.watch_html)
         self.vid_info_url = extract.video_info_url(

pytube/cipher.py

@@ -14,28 +14,27 @@ signature and decoding it.
 
 """
 
-import logging
-import pprint
 import re
 from itertools import chain
+from typing import List, Tuple, Dict, Callable
 
 from pytube.exceptions import RegexMatchError
-from pytube.helpers import regex_search
+from pytube.helpers import regex_search, create_logger
 
+logger = create_logger()
 
-logger = logging.getLogger(__name__)
 
-
-def get_initial_function_name(js):
+def get_initial_function_name(js: str) -> str:
     """Extract the name of the function responsible for computing the signature.
 
     :param str js:
         The contents of the base.js asset file.
-
+    :rtype: str
+    :returns:
+       Function name from regex match
     """
-    # c&&d.set("signature", EE(c));
 
-    pattern = [
+    function_patterns = [
         r"\b[cs]\s*&&\s*[adf]\.set\([^,]+\s*,\s*encodeURIComponent\s*\(\s*(?P<sig>[a-zA-Z0-9$]+)\(",  # noqa: E501
         r"\b[a-zA-Z0-9]+\s*&&\s*[a-zA-Z0-9]+\.set\([^,]+\s*,\s*encodeURIComponent\s*\(\s*(?P<sig>[a-zA-Z0-9$]+)\(",  # noqa: E501
         r'(?P<sig>[a-zA-Z0-9$]+)\s*=\s*function\(\s*a\s*\)\s*{\s*a\s*=\s*a\.split\(\s*""\s*\)',  # noqa: E501
@@ -50,10 +49,19 @@ def get_initial_function_name(js):
     ]
 
     logger.debug("finding initial function name")
-    return regex_search(pattern, js, group=1)
+    for pattern in function_patterns:
+        regex = re.compile(pattern)
+        results = regex.search(js)
+        if results:
+            logger.debug(
+                "finished regex search, matched: {pattern}".format(pattern=pattern)
+            )
+            return results.group(1)
+
+    raise RegexMatchError(caller="get_initial_function_name", pattern="multiple")
 
 
-def get_transform_plan(js):
+def get_transform_plan(js: str) -> List[str]:
     """Extract the "transform plan".
 
     The "transform plan" is the functions that the ciphered signature is
@@ -80,7 +88,7 @@ def get_transform_plan(js):
     return regex_search(pattern, js, group=1).split(";")
 
 
-def get_transform_object(js, var):
+def get_transform_object(js: str, var: str) -> List[str]:
     """Extract the "transform object".
 
     The "transform object" contains the function definitions referenced in the
@@ -104,14 +112,15 @@ def get_transform_object(js, var):
     """
     pattern = r"var %s={(.*?)};" % re.escape(var)
     logger.debug("getting transform object")
-    return (
-        regex_search(pattern, js, group=1, flags=re.DOTALL)
-        .replace("\n", " ")
-        .split(", ")
-    )
+    regex = re.compile(pattern, flags=re.DOTALL)
+    results = regex.search(js)
+    if not results:
+        raise RegexMatchError(caller="get_transform_object", pattern=pattern)
+
+    return results.group(1).replace("\n", " ").split(", ")
 
 
-def get_transform_map(js, var):
+def get_transform_map(js: str, var: str) -> Dict:
     """Build a transform function lookup.
 
     Build a lookup table of obfuscated JavaScript function names to the
@@ -189,7 +198,7 @@ def swap(arr, b):
     return list(chain([arr[r]], arr[1:r], [arr[0]], arr[r + 1 :]))
 
 
-def map_functions(js_func):
+def map_functions(js_func: str) -> Callable:
     """For a given JavaScript transform function, return the Python equivalent.
 
     :param str js_func:
@@ -213,12 +222,10 @@ def map_functions(js_func):
     for pattern, fn in mapper:
         if re.search(pattern, js_func):
             return fn
-    raise RegexMatchError(
-        "could not find python equivalent function for: ", js_func,
-    )
+    raise RegexMatchError(caller="map_functions", pattern="multiple")
 
 
-def parse_function(js_func):
+def parse_function(js_func: str) -> Tuple[str, int]:
     """Parse the Javascript transform function.
 
     Break a JavaScript transform function down into a two element ``tuple``
@@ -237,7 +244,13 @@ def parse_function(js_func):
 
     """
     logger.debug("parsing transform function")
-    return regex_search(r"\w+\.(\w+)\(\w,(\d+)\)", js_func, groups=True)
+    pattern = r"\w+\.(\w+)\(\w,(\d+)\)"
+    regex = re.compile(pattern)
+    results = regex.search(js_func)
+    if not results:
+        raise RegexMatchError(caller="parse_function", pattern=pattern)
+    fn_name, fn_arg = results.groups()
+    return fn_name, int(fn_arg)
 
 
 def get_signature(js: str, ciphered_signature: str) -> str:
@@ -255,24 +268,23 @@ def get_signature(js: str, ciphered_signature: str) -> str:
 
     """
     transform_plan = get_transform_plan(js)
-    # DE.AJ(a,15) => DE, AJ(a,15)
     var, _ = transform_plan[0].split(".")
     transform_map = get_transform_map(js, var)
     signature = [s for s in ciphered_signature]
 
     for js_func in transform_plan:
         name, argument = parse_function(js_func)
-        signature = transform_map[name](signature, int(argument))
+        signature = transform_map[name](signature, argument)
         logger.debug(
-            "applied transform function\n%s",
-            pprint.pformat(
-                {
-                    "output": "".join(signature),
-                    "js_function": name,
-                    "argument": int(argument),
-                    "function": transform_map[name],
-                },
-                indent=2,
-            ),
+            "applied transform function\n"
+            "output: %s\n"
+            "js_function: %s\n"
+            "argument: %d\n"
+            "function: %s",
+            "".join(signature),
+            name,
+            argument,
+            transform_map[name],
         )
+
     return "".join(signature)

pytube/exceptions.py

@@ -1,6 +1,6 @@
 # -*- coding: utf-8 -*-
 """Library specific exception definitions."""
-import sys
+from typing import Union, Pattern
 
 
 class PytubeError(Exception):
@@ -15,26 +15,25 @@ class PytubeError(Exception):
 class ExtractError(PytubeError):
     """Data extraction based exception."""
 
-    def __init__(self, msg: str, video_id: str = "unknown id"):
-        """Construct an instance of a :class:`ExtractError <ExtractError>`.
-
-        :param str msg:
-            User defined error message.
-        :param str video_id:
-            A YouTube video identifier.
-        """
-        if video_id is not None:
-            msg = "{video_id}: {msg}".format(video_id=video_id, msg=msg)
-
-        super(ExtractError, self).__init__(msg)
-
-        self.exc_info = sys.exc_info()
-        self.video_id = video_id
-
 
 class RegexMatchError(ExtractError):
     """Regex pattern did not return any matches."""
 
+    def __init__(self, caller: str, pattern: Union[str, Pattern]):
+        """
+        :param str caller:
+            Calling function
+        :param str pattern:
+            Pattern that failed to match
+        """
+        super().__init__(
+            "{caller}: could not find match for {pattern}".format(
+                caller=caller, pattern=pattern
+            )
+        )
+        self.caller = caller
+        self.pattern = pattern
+
 
 class LiveStreamError(ExtractError):
     """Video is a live stream."""
@@ -43,6 +42,15 @@ class LiveStreamError(ExtractError):
 class VideoUnavailable(PytubeError):
     """Video is unavailable."""
 
+    def __init__(self, video_id: str):
+        """
+        :param str video_id:
+            A YouTube video identifier.
+        """
+        super().__init__("{video_id} is unavailable".format(video_id=video_id))
+
+        self.video_id = video_id
+
 
 class HTMLParseError(PytubeError):
     """HTML could not be parsed"""

pytube/extract.py

@@ -1,6 +1,7 @@
 # -*- coding: utf-8 -*-
 """This module contains all non-cipher related data extraction logic."""
 import json
+import re
 from collections import OrderedDict
 
 from html.parser import HTMLParser
@@ -170,7 +171,11 @@ def mime_type_codec(mime_type_codec: str) -> Tuple[str, List[str]]:
 
     """
     pattern = r"(\w+\/\w+)\;\scodecs=\"([a-zA-Z-0-9.,\s]*)\""
-    mime_type, codecs = regex_search(pattern, mime_type_codec, groups=True)
+    regex = re.compile(pattern)
+    results = regex.search(mime_type_codec)
+    if not results:
+        raise RegexMatchError(caller="mime_type_codec", pattern=pattern)
+    mime_type, codecs = results.groups()
     return mime_type, [c.strip() for c in codecs.split(",")]
 
 

pytube/helpers.py

@@ -7,69 +7,34 @@ import re
 
 from pytube.exceptions import RegexMatchError
 
-
 logger = logging.getLogger(__name__)
 
 
-def regex_search(pattern, string, groups=False, group=None, flags=0):
+def regex_search(pattern: str, string: str, group: int) -> str:
     """Shortcut method to search a string for a given pattern.
 
     :param str pattern:
         A regular expression pattern.
     :param str string:
         A target string to search.
-    :param bool groups:
-        Should the return value be ``.groups()``.
     :param int group:
         Index of group to return.
-    :param int flags:
-        Expression behavior modifiers.
     :rtype:
         str or tuple
     :returns:
         Substring pattern matches.
     """
-    if type(pattern) == list:
-        for p in pattern:
-            regex = re.compile(p, flags)
-            results = regex.search(string)
-            if not results:
-                raise RegexMatchError(
-                    "regex pattern ({pattern}) had zero matches".format(pattern=p),
-                )
-            else:
-                logger.debug(
-                    "finished regex search: %s",
-                    pprint.pformat(
-                        {"pattern": p, "results": results.group(0),}, indent=2,
-                    ),
-                )
-                if groups:
-                    return results.groups()
-                elif group is not None:
-                    return results.group(group)
-                else:
-                    return results
-    else:
-        regex = re.compile(pattern, flags)
-        results = regex.search(string)
-        if not results:
-            raise RegexMatchError(
-                "regex pattern ({pattern}) had zero matches".format(pattern=pattern),
-            )
-        else:
-            logger.debug(
-                "finished regex search: %s",
-                pprint.pformat(
-                    {"pattern": pattern, "results": results.group(0),}, indent=2,
-                ),
-            )
-            if groups:
-                return results.groups()
-            elif group is not None:
-                return results.group(group)
-            else:
-                return results
+    regex = re.compile(pattern)
+    results = regex.search(string)
+    if not results:
+        raise RegexMatchError(caller="regex_search", pattern=pattern)
+
+    logger.debug(
+        "finished regex search: %s",
+        pprint.pformat({"pattern": pattern, "results": results.group(0),}, indent=2,),
+    )
+
+    return results.group(group)
 
 
 def safe_filename(s: str, max_length: int = 255) -> str:

pytube/request.py

@@ -3,12 +3,8 @@
 from urllib.request import Request
 from urllib.request import urlopen
 
-# 403 forbidden fix
 
-
-def get(
-    url, headers=False, streaming=False, chunk_size=8 * 1024,
-):
+def get(url, headers=False, streaming=False, chunk_size=8192):
     """Send an http GET request.
 
     :param str url:
@@ -18,7 +14,7 @@ def get(
     :param bool streaming:
         Returns the response body in chunks via a generator.
     :param int chunk_size:
-        The size in bytes of each chunk.
+        The size in bytes of each chunk. Defaults to 8*1024
     """
 
     req = Request(url, headers={"User-Agent": "Mozilla/5.0"})

pytube/streams.py

@@ -236,21 +236,20 @@ class Stream:
                 prefix=safe_filename(filename_prefix), filename=filename,
             )
 
-        # file path
-        fp = os.path.join(output_path, filename)
+        file_path = os.path.join(output_path, filename)
         bytes_remaining = self.filesize
         logger.debug(
-            "downloading (%s total bytes) file to %s", self.filesize, fp,
+            "downloading (%s total bytes) file to %s", self.filesize, file_path,
         )
 
-        with open(fp, "wb") as fh:
+        with open(file_path, "wb") as fh:
             for chunk in request.get(self.url, streaming=True):
                 # reduce the (bytes) remainder by the length of the chunk.
                 bytes_remaining -= len(chunk)
                 # send to the on_progress callback.
                 self.on_progress(chunk, fh, bytes_remaining)
         self.on_complete(fh)
-        return fp
+        return file_path
 
     def stream_to_buffer(self) -> io.BytesIO:
         """Write the media stream to buffer

tests/test_captions.py

@@ -0,0 +1,41 @@
+from pytube import Caption, CaptionQuery
+
+
+def test_float_to_srt_time_format():
+    caption1 = Caption(
+        {"url": "url1", "name": {"simpleText": "name1"}, "languageCode": "en"}
+    )
+    assert caption1.float_to_srt_time_format(3.89) == "00:00:03,890"
+
+
+def test_caption_query_all():
+    caption1 = Caption(
+        {"url": "url1", "name": {"simpleText": "name1"}, "languageCode": "en"}
+    )
+    caption2 = Caption(
+        {"url": "url2", "name": {"simpleText": "name2"}, "languageCode": "fr"}
+    )
+    caption_query = CaptionQuery(captions=[caption1, caption2])
+    assert caption_query.captions == [caption1, caption2]
+
+
+def test_caption_query_get_by_language_code_when_exists():
+    caption1 = Caption(
+        {"url": "url1", "name": {"simpleText": "name1"}, "languageCode": "en"}
+    )
+    caption2 = Caption(
+        {"url": "url2", "name": {"simpleText": "name2"}, "languageCode": "fr"}
+    )
+    caption_query = CaptionQuery(captions=[caption1, caption2])
+    assert caption_query.get_by_language_code("en") == caption1
+
+
+def test_caption_query_get_by_language_code_when_not_exists():
+    caption1 = Caption(
+        {"url": "url1", "name": {"simpleText": "name1"}, "languageCode": "en"}
+    )
+    caption2 = Caption(
+        {"url": "url2", "name": {"simpleText": "name2"}, "languageCode": "fr"}
+    )
+    caption_query = CaptionQuery(captions=[caption1, caption2])
+    assert caption_query.get_by_language_code("hello") is None

tests/test_cipher.py

@@ -8,3 +8,29 @@ from pytube.exceptions import RegexMatchError
 def test_map_functions():
     with pytest.raises(RegexMatchError):
         cipher.map_functions("asdf")
+
+
+def test_get_initial_function_name_with_no_match_should_error():
+    with pytest.raises(RegexMatchError):
+        cipher.get_initial_function_name("asdf")
+
+
+def test_get_transform_object_with_no_match_should_error():
+    with pytest.raises(RegexMatchError):
+        cipher.get_transform_object("asdf", var="lt")
+
+
+def test_parse_function_with_match():
+    fn_name, fn_arg = cipher.parse_function("DE.AJ(a,15)")
+    assert fn_name == "AJ"
+    assert fn_arg == 15
+
+
+def test_parse_function_with_no_match_should_error():
+    with pytest.raises(RegexMatchError):
+        cipher.parse_function("asdf")
+
+
+def test_reverse():
+    reversed_array = cipher.reverse([1, 2, 3, 4], None)
+    assert reversed_array == [4, 3, 2, 1]

tests/test_exceptions.py

@@ -1,9 +1,17 @@
 # -*- coding: utf-8 -*-
-from pytube.exceptions import ExtractError
+from pytube.exceptions import VideoUnavailable, RegexMatchError
 
 
-def test_is_expected():
+def test_video_unavailable():
     try:
-        raise ExtractError("ppfff", video_id="YLnZklYFe7E")
-    except ExtractError as e:
+        raise VideoUnavailable(video_id="YLnZklYFe7E")
+    except VideoUnavailable as e:
         assert e.video_id == "YLnZklYFe7E"
+        assert str(e) == "YLnZklYFe7E is unavailable"
+
+
+def test_regex_match_error():
+    try:
+        raise RegexMatchError(caller="hello", pattern="*")
+    except RegexMatchError as e:
+        assert str(e) == "hello: could not find match for *"

tests/test_extract.py

@@ -1,5 +1,9 @@
 # -*- coding: utf-8 -*-
 """Unit tests for the :module:`extract <extract>` module."""
+import pytest
+
+from pytube.exceptions import RegexMatchError
+
 from pytube import extract
 
 
@@ -61,3 +65,19 @@ def test_get_vid_desc(cipher_signature):
         "http://weibo.com/psyoppa"
     )
     assert extract.get_vid_descr(cipher_signature.watch_html) == expected
+
+
+def test_eurl():
+    url = extract.eurl("videoid")
+    assert url == "https://youtube.googleapis.com/v/videoid"
+
+
+def test_mime_type_codec():
+    mime_type, mime_subtype = extract.mime_type_codec('audio/webm; codecs="opus"')
+    assert mime_type == "audio/webm"
+    assert mime_subtype == ["opus"]
+
+
+def test_mime_type_codec_with_no_match_should_error():
+    with pytest.raises(RegexMatchError):
+        extract.mime_type_codec("audio/webm")

tests/test_helpers.py

@@ -7,12 +7,11 @@ from pytube.exceptions import RegexMatchError
 
 def test_regex_search_no_match():
     with pytest.raises(RegexMatchError):
-        helpers.regex_search("^a$", "", groups=True)
+        helpers.regex_search("^a$", "", group=0)
 
 
 def test_regex_search():
-    # TODO(nficano): should check isinstance
-    assert helpers.regex_search("^a$", "a") is not None
+    assert helpers.regex_search("^a$", "a", group=0) == "a"
 
 
 def test_safe_filename():

tests/test_main.py

@@ -1,7 +1,10 @@
 # -*- coding: utf-8 -*-
 from unittest import mock
 
+import pytest
+
 from pytube import YouTube
+from pytube.exceptions import VideoUnavailable
 
 
 @mock.patch("pytube.__main__.YouTube")
@@ -10,3 +13,24 @@ def test_prefetch_deferred(MockYouTube):
     instance.prefetch_descramble.return_value = None
     YouTube("https://www.youtube.com/watch?v=9bZkp7q19f0", True)
     assert not instance.prefetch_descramble.called
+
+
+@mock.patch("urllib.request.install_opener")
+def test_install_proxy(opener):
+    proxies = {"http": "http://www.example.com:3128/"}
+    YouTube(
+        "https://www.youtube.com/watch?v=9bZkp7q19f0",
+        defer_prefetch_init=True,
+        proxies=proxies,
+    )
+    opener.assert_called()
+
+
+@mock.patch("pytube.request.get")
+def test_video_unavailable(get):
+    get.return_value = None
+    youtube = YouTube(
+        "https://www.youtube.com/watch?v=9bZkp7q19f0", defer_prefetch_init=True
+    )
+    with pytest.raises(VideoUnavailable):
+        youtube.prefetch()

tests/test_streams.py

@@ -3,8 +3,6 @@ import random
 
 from unittest import mock
 
-import pytest
-
 from pytube import request
 from pytube import Stream
 
@@ -107,40 +105,53 @@ def test_author(cipher_signature):
     assert cipher_signature.author == expected
 
 
-@pytest.mark.skip
+def test_thumbnail_when_in_details(cipher_signature):
+    expected = "some url"
+    cipher_signature.player_config_args = {
+        "player_response": {
+            "videoDetails": {"thumbnail": {"thumbnails": [{"url": expected}]}}
+        }
+    }
+    assert cipher_signature.thumbnail_url == expected
+
+
+def test_thumbnail_when_not_in_details(cipher_signature):
+    expected = "https://img.youtube.com/vi/9bZkp7q19f0/maxresdefault.jpg"
+    cipher_signature.player_config_args = {}
+    assert cipher_signature.thumbnail_url == expected
+
+
 def test_repr_for_audio_streams(cipher_signature):
     stream = str(cipher_signature.streams.filter(only_audio=True).first())
     expected = (
-        '<Stream: itag="140" mime_type="audio/mp4" abr="128kbps" ' 'acodec="mp4a.40.2">'
+        '<Stream: itag="140" mime_type="audio/mp4" abr="128kbps" '
+        'acodec="mp4a.40.2" progressive="False" type="audio">'
     )
     assert stream == expected
 
 
-@pytest.mark.skip
 def test_repr_for_video_streams(cipher_signature):
     stream = str(cipher_signature.streams.filter(only_video=True).first())
     expected = (
-        '<Stream: itag="137" mime_type="video/mp4" res="1080p" '
-        'fps="30fps" vcodec="avc1.640028">'
+        '<Stream: itag="137" mime_type="video/mp4" res="1080p" fps="30fps" '
+        'vcodec="avc1.640028" progressive="False" type="video">'
     )
     assert stream == expected
 
 
-@pytest.mark.skip
 def test_repr_for_progressive_streams(cipher_signature):
     stream = str(cipher_signature.streams.filter(progressive=True).first())
     expected = (
-        '<Stream: itag="18" mime_type="video/mp4" res="360p" '
-        'fps="30fps" vcodec="avc1.42001E" acodec="mp4a.40.2">'
+        '<Stream: itag="18" mime_type="video/mp4" res="360p" fps="30fps" '
+        'vcodec="avc1.42001E" acodec="mp4a.40.2" progressive="True" type="video">'
     )
     assert stream == expected
 
 
-@pytest.mark.skip
 def test_repr_for_adaptive_streams(cipher_signature):
     stream = str(cipher_signature.streams.filter(adaptive=True).first())
     expected = (
-        '<Stream: itag="137" mime_type="video/mp4" res="1080p" '
-        'fps="30fps" vcodec="avc1.640028">'
+        '<Stream: itag="137" mime_type="video/mp4" res="1080p" fps="30fps" '
+        'vcodec="avc1.640028" progressive="False" type="video">'
     )
     assert stream == expected