Source code for pytube.helpers

"""Various helper functions implemented by pytube."""
import functools
import gzip
import json
import logging
import os
import re
import warnings
from typing import Any, Callable, Dict, List, Optional, TypeVar
from urllib import request

from pytube.exceptions import RegexMatchError

logger = logging.getLogger(__name__)

[docs]class DeferredGeneratorList: """A wrapper class for deferring list generation. Pytube has some continuation generators that create web calls, which means that any time a full list is requested, all of those web calls must be made at once, which could lead to slowdowns. This will allow individual elements to be queried, so that slowdowns only happen as necessary. For example, you can iterate over elements in the list without accessing them all simultaneously. This should allow for speed improvements for playlist and channel interactions. """ def __init__(self, generator): """Construct a :class:`DeferredGeneratorList <DeferredGeneratorList>`. :param generator generator: The deferrable generator to create a wrapper for. :param func func: (Optional) A function to call on the generator items to produce the list. """ self.gen = generator self._elements = [] def __eq__(self, other): """We want to mimic list behavior for comparison.""" return list(self) == other def __getitem__(self, key) -> Any: """Only generate items as they're asked for.""" # We only allow querying with indexes. if not isinstance(key, (int, slice)): raise TypeError('Key must be either a slice or int.') # Convert int keys to slice key_slice = key if isinstance(key, int): key_slice = slice(key, key + 1, 1) # Generate all elements up to the final item while len(self._elements) < key_slice.stop: try: next_item = next(self.gen) except StopIteration: # If we can't find enough elements for the slice, raise an IndexError raise IndexError else: self._elements.append(next_item) return self._elements[key] def __iter__(self): """Custom iterator for dynamically generated list.""" iter_index = 0 while True: try: curr_item = self[iter_index] except IndexError: return else: yield curr_item iter_index += 1 def __next__(self) -> Any: """Fetch next element in iterator.""" try: curr_element = self[self.iter_index] except IndexError: raise StopIteration self.iter_index += 1 return curr_element # noqa:R504 def __len__(self) -> int: """Return length of list of all items.""" self.generate_all() return len(self._elements) def __repr__(self) -> str: """String representation of all items.""" self.generate_all() return str(self._elements) def __reversed__(self): self.generate_all() return self._elements[::-1]
[docs] def generate_all(self): """Generate all items.""" while True: try: next_item = next(self.gen) except StopIteration: break else: self._elements.append(next_item)
[docs]def safe_filename(s: str, max_length: int = 255) -> str: """Sanitize a string making it safe to use as a filename. This function was based off the limitations outlined here: :param str s: A string to make safe for use as a file name. :param int max_length: The maximum filename character length. :rtype: str :returns: A sanitized string. """ # Characters in range 0-31 (0x00-0x1F) are not allowed in ntfs filenames. ntfs_characters = [chr(i) for i in range(0, 31)] characters = [ r'"', r"\#", r"\$", r"\%", r"'", r"\*", r"\,", r"\.", r"\/", r"\:", r'"', r"\;", r"\<", r"\>", r"\?", r"\\", r"\^", r"\|", r"\~", r"\\\\", ] pattern = "|".join(ntfs_characters + characters) regex = re.compile(pattern, re.UNICODE) filename = regex.sub("", s) return filename[:max_length].rsplit(" ", 0)[0]
[docs]def setup_logger(level: int = logging.ERROR, log_filename: Optional[str] = None) -> None: """Create a configured instance of logger. :param int level: Describe the severity level of the logs to handle. """ fmt = "[%(asctime)s] %(levelname)s in %(module)s: %(message)s" date_fmt = "%H:%M:%S" formatter = logging.Formatter(fmt, datefmt=date_fmt) # logger = logging.getLogger("pytube") logger.setLevel(level) stream_handler = logging.StreamHandler() stream_handler.setFormatter(formatter) logger.addHandler(stream_handler) if log_filename is not None: file_handler = logging.FileHandler(log_filename) file_handler.setFormatter(formatter) logger.addHandler(file_handler)
GenericType = TypeVar("GenericType")
[docs]def cache(func: Callable[..., GenericType]) -> GenericType: """ mypy compatible annotation wrapper for lru_cache""" return functools.lru_cache()(func) # type: ignore
[docs]def deprecated(reason: str) -> Callable: """ This is a decorator which can be used to mark functions as deprecated. It will result in a warning being emitted when the function is used. """ def decorator(func1): message = "Call to deprecated function {name} ({reason})." @functools.wraps(func1) def new_func1(*args, **kwargs): warnings.simplefilter("always", DeprecationWarning) warnings.warn( message.format(name=func1.__name__, reason=reason), category=DeprecationWarning, stacklevel=2, ) warnings.simplefilter("default", DeprecationWarning) return func1(*args, **kwargs) return new_func1 return decorator
[docs]def target_directory(output_path: Optional[str] = None) -> str: """ Function for determining target directory of a download. Returns an absolute path (if relative one given) or the current path (if none given). Makes directory if it does not exist. :type output_path: str :rtype: str :returns: An absolute directory path as a string. """ if output_path: if not os.path.isabs(output_path): output_path = os.path.join(os.getcwd(), output_path) else: output_path = os.getcwd() os.makedirs(output_path, exist_ok=True) return output_path
def install_proxy(proxy_handler: Dict[str, str]) -> None: proxy_support = request.ProxyHandler(proxy_handler) opener = request.build_opener(proxy_support) request.install_opener(opener)
[docs]def uniqueify(duped_list: List) -> List: """Remove duplicate items from a list, while maintaining list order. :param List duped_list List to remove duplicates from :return List result De-duplicated list """ seen: Dict[Any, bool] = {} result = [] for item in duped_list: if item in seen: continue seen[item] = True result.append(item) return result
[docs]def generate_all_html_json_mocks(): """Regenerate the video mock json files for all current test videos. This should automatically output to the test/mocks directory. """ test_vid_ids = [ '2lAe1cqCOXo', '5YceQ8YqYMc', 'irauhITDrsE', 'm8uHb5jIGN8', 'QRS8MkLhQmM', 'WXxV9g7lsFE' ] for vid_id in test_vid_ids: create_mock_html_json(vid_id)
[docs]def create_mock_html_json(vid_id) -> Dict[str, Any]: """Generate a json.gz file with sample html responses. :param str vid_id YouTube video id :return dict data Dict used to generate the json.gz file """ from pytube import YouTube gzip_filename = 'yt-video-%s-html.json.gz' % vid_id # Get the pytube directory in order to navigate to /tests/mocks pytube_dir_path = os.path.abspath( os.path.join( os.path.dirname(__file__), os.path.pardir ) ) pytube_mocks_path = os.path.join(pytube_dir_path, 'tests', 'mocks') gzip_filepath = os.path.join(pytube_mocks_path, gzip_filename) yt = YouTube(f'{vid_id}') html_data = { 'url': yt.watch_url, 'js': yt.js, 'embed_html': yt.embed_html, 'watch_html': yt.watch_html, 'vid_info': yt.vid_info }'Outputing json.gz file to {gzip_filepath}') with, 'wb') as f: f.write(json.dumps(html_data).encode('utf-8')) return html_data