mirror of
https://github.com/yt-dlp/yt-dlp.git
synced 2024-11-02 06:27:21 +00:00
[utils] Allow partial application for more functions (#11391)
Also adds the `trim_str` traversal helper Authored by: bashonly, Grub4K Co-authored-by: Simon Sawicki <contact@grub4k.xyz>
This commit is contained in:
parent
76802f4613
commit
b6dc2c49e8
|
@ -12,9 +12,10 @@
|
|||
str_or_none,
|
||||
)
|
||||
from yt_dlp.utils.traversal import (
|
||||
traverse_obj,
|
||||
require,
|
||||
subs_list_to_dict,
|
||||
traverse_obj,
|
||||
trim_str,
|
||||
)
|
||||
|
||||
_TEST_DATA = {
|
||||
|
@ -495,6 +496,20 @@ def test_subs_list_to_dict(self):
|
|||
{'url': 'https://example.com/subs/en2', 'ext': 'ext'},
|
||||
]}, '`quality` key should sort subtitle list accordingly'
|
||||
|
||||
def test_trim_str(self):
|
||||
with pytest.raises(TypeError):
|
||||
trim_str('positional')
|
||||
|
||||
assert callable(trim_str(start='a'))
|
||||
assert trim_str(start='ab')('abc') == 'c'
|
||||
assert trim_str(end='bc')('abc') == 'a'
|
||||
assert trim_str(start='a', end='c')('abc') == 'b'
|
||||
assert trim_str(start='ab', end='c')('abc') == ''
|
||||
assert trim_str(start='a', end='bc')('abc') == ''
|
||||
assert trim_str(start='ab', end='bc')('abc') == ''
|
||||
assert trim_str(start='abc', end='abc')('abc') == ''
|
||||
assert trim_str(start='', end='')('abc') == 'abc'
|
||||
|
||||
|
||||
class TestDictGet:
|
||||
def test_dict_get(self):
|
||||
|
|
|
@ -4,6 +4,7 @@
|
|||
import os
|
||||
import sys
|
||||
import unittest
|
||||
import unittest.mock
|
||||
import warnings
|
||||
import datetime as dt
|
||||
|
||||
|
@ -71,6 +72,7 @@
|
|||
intlist_to_bytes,
|
||||
iri_to_uri,
|
||||
is_html,
|
||||
join_nonempty,
|
||||
js_to_json,
|
||||
limit_length,
|
||||
locked_file,
|
||||
|
@ -343,11 +345,13 @@ def test_remove_start(self):
|
|||
self.assertEqual(remove_start(None, 'A - '), None)
|
||||
self.assertEqual(remove_start('A - B', 'A - '), 'B')
|
||||
self.assertEqual(remove_start('B - A', 'A - '), 'B - A')
|
||||
self.assertEqual(remove_start('non-empty', ''), 'non-empty')
|
||||
|
||||
def test_remove_end(self):
|
||||
self.assertEqual(remove_end(None, ' - B'), None)
|
||||
self.assertEqual(remove_end('A - B', ' - B'), 'A')
|
||||
self.assertEqual(remove_end('B - A', ' - B'), 'B - A')
|
||||
self.assertEqual(remove_end('non-empty', ''), 'non-empty')
|
||||
|
||||
def test_remove_quotes(self):
|
||||
self.assertEqual(remove_quotes(None), None)
|
||||
|
@ -2148,6 +2152,16 @@ def run_shell(args):
|
|||
assert run_shell(args) == expected
|
||||
assert run_shell(shell_quote(args, shell=True)) == expected
|
||||
|
||||
def test_partial_application(self):
|
||||
assert callable(int_or_none(scale=10)), 'missing positional parameter should apply partially'
|
||||
assert int_or_none(10, scale=0.1) == 100, 'positionally passed argument should call function'
|
||||
assert int_or_none(v=10) == 10, 'keyword passed positional should call function'
|
||||
assert int_or_none(scale=0.1)(10) == 100, 'call after partial applicatino should call the function'
|
||||
|
||||
assert callable(join_nonempty(delim=', ')), 'varargs positional should apply partially'
|
||||
assert callable(join_nonempty()), 'varargs positional should apply partially'
|
||||
assert join_nonempty(None, delim=', ') == '', 'passed varargs should call the function'
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
unittest.main()
|
||||
|
|
|
@ -212,6 +212,23 @@ def write_json_file(obj, fn):
|
|||
raise
|
||||
|
||||
|
||||
def partial_application(func):
|
||||
sig = inspect.signature(func)
|
||||
required_args = [
|
||||
param.name for param in sig.parameters.values()
|
||||
if param.kind in (inspect.Parameter.POSITIONAL_ONLY, inspect.Parameter.POSITIONAL_OR_KEYWORD, inspect.Parameter.VAR_POSITIONAL)
|
||||
if param.default is inspect.Parameter.empty
|
||||
]
|
||||
|
||||
@functools.wraps(func)
|
||||
def wrapped(*args, **kwargs):
|
||||
if set(required_args[len(args):]).difference(kwargs):
|
||||
return functools.partial(func, *args, **kwargs)
|
||||
return func(*args, **kwargs)
|
||||
|
||||
return wrapped
|
||||
|
||||
|
||||
def find_xpath_attr(node, xpath, key, val=None):
|
||||
""" Find the xpath xpath[@key=val] """
|
||||
assert re.match(r'^[a-zA-Z_-]+$', key)
|
||||
|
@ -1192,6 +1209,7 @@ def extract_timezone(date_str, default=None):
|
|||
return timezone, date_str
|
||||
|
||||
|
||||
@partial_application
|
||||
def parse_iso8601(date_str, delimiter='T', timezone=None):
|
||||
""" Return a UNIX timestamp from the given date """
|
||||
|
||||
|
@ -1269,6 +1287,7 @@ def unified_timestamp(date_str, day_first=True):
|
|||
return calendar.timegm(timetuple) + pm_delta * 3600 - timezone.total_seconds()
|
||||
|
||||
|
||||
@partial_application
|
||||
def determine_ext(url, default_ext='unknown_video'):
|
||||
if url is None or '.' not in url:
|
||||
return default_ext
|
||||
|
@ -1944,7 +1963,7 @@ def remove_start(s, start):
|
|||
|
||||
|
||||
def remove_end(s, end):
|
||||
return s[:-len(end)] if s is not None and s.endswith(end) else s
|
||||
return s[:-len(end)] if s is not None and end and s.endswith(end) else s
|
||||
|
||||
|
||||
def remove_quotes(s):
|
||||
|
@ -1973,6 +1992,7 @@ def base_url(url):
|
|||
return re.match(r'https?://[^?#]+/', url).group()
|
||||
|
||||
|
||||
@partial_application
|
||||
def urljoin(base, path):
|
||||
if isinstance(path, bytes):
|
||||
path = path.decode()
|
||||
|
@ -1988,21 +2008,6 @@ def urljoin(base, path):
|
|||
return urllib.parse.urljoin(base, path)
|
||||
|
||||
|
||||
def partial_application(func):
|
||||
sig = inspect.signature(func)
|
||||
|
||||
@functools.wraps(func)
|
||||
def wrapped(*args, **kwargs):
|
||||
try:
|
||||
sig.bind(*args, **kwargs)
|
||||
except TypeError:
|
||||
return functools.partial(func, *args, **kwargs)
|
||||
else:
|
||||
return func(*args, **kwargs)
|
||||
|
||||
return wrapped
|
||||
|
||||
|
||||
@partial_application
|
||||
def int_or_none(v, scale=1, default=None, get_attr=None, invscale=1, base=None):
|
||||
if get_attr and v is not None:
|
||||
|
@ -2583,6 +2588,7 @@ def urlencode_postdata(*args, **kargs):
|
|||
return urllib.parse.urlencode(*args, **kargs).encode('ascii')
|
||||
|
||||
|
||||
@partial_application
|
||||
def update_url(url, *, query_update=None, **kwargs):
|
||||
"""Replace URL components specified by kwargs
|
||||
@param url str or parse url tuple
|
||||
|
@ -2603,6 +2609,7 @@ def update_url(url, *, query_update=None, **kwargs):
|
|||
return urllib.parse.urlunparse(url._replace(**kwargs))
|
||||
|
||||
|
||||
@partial_application
|
||||
def update_url_query(url, query):
|
||||
return update_url(url, query_update=query)
|
||||
|
||||
|
@ -2924,6 +2931,7 @@ def error_to_str(err):
|
|||
return f'{type(err).__name__}: {err}'
|
||||
|
||||
|
||||
@partial_application
|
||||
def mimetype2ext(mt, default=NO_DEFAULT):
|
||||
if not isinstance(mt, str):
|
||||
if default is not NO_DEFAULT:
|
||||
|
@ -4664,6 +4672,7 @@ def to_high_limit_path(path):
|
|||
return path
|
||||
|
||||
|
||||
@partial_application
|
||||
def format_field(obj, field=None, template='%s', ignore=NO_DEFAULT, default='', func=IDENTITY):
|
||||
val = traversal.traverse_obj(obj, *variadic(field))
|
||||
if not val if ignore is NO_DEFAULT else val in variadic(ignore):
|
||||
|
@ -4828,6 +4837,7 @@ def number_of_digits(number):
|
|||
return len('%d' % number)
|
||||
|
||||
|
||||
@partial_application
|
||||
def join_nonempty(*values, delim='-', from_dict=None):
|
||||
if from_dict is not None:
|
||||
values = (traversal.traverse_obj(from_dict, variadic(v)) for v in values)
|
||||
|
@ -5278,6 +5288,7 @@ def report_retry(e, count, retries, *, sleep_func, info, warn, error=None, suffi
|
|||
time.sleep(delay)
|
||||
|
||||
|
||||
@partial_application
|
||||
def make_archive_id(ie, video_id):
|
||||
ie_key = ie if isinstance(ie, str) else ie.ie_key()
|
||||
return f'{ie_key.lower()} {video_id}'
|
||||
|
|
|
@ -435,6 +435,20 @@ def find_elements(*, tag=None, cls=None, attr=None, value=None, html=False):
|
|||
return functools.partial(func, cls)
|
||||
|
||||
|
||||
def trim_str(*, start=None, end=None):
|
||||
def trim(s):
|
||||
if s is None:
|
||||
return None
|
||||
start_idx = 0
|
||||
if start and s.startswith(start):
|
||||
start_idx = len(start)
|
||||
if end and s.endswith(end):
|
||||
return s[start_idx:-len(end)]
|
||||
return s[start_idx:]
|
||||
|
||||
return trim
|
||||
|
||||
|
||||
def get_first(obj, *paths, **kwargs):
|
||||
return traverse_obj(obj, *((..., *variadic(keys)) for keys in paths), **kwargs, get_all=False)
|
||||
|
||||
|
|
Loading…
Reference in a new issue