Browse Source

[utils] Support `filter` traversal key

Thx yt-dlp/yt-dlp#10653
dirkf 1 month ago
parent
commit
96419fa706
3 changed files with 22 additions and 0 deletions
  1. 8 0
      test/test_traversal.py
  2. 6 0
      youtube_dl/compat.py
  3. 8 0
      youtube_dl/utils.py

+ 8 - 0
test/test_traversal.py

@@ -473,6 +473,14 @@ class TestTraversal(_TestCase):
         self.assertIs(traverse_obj(morsel, [(None,), any]), morsel,
                       msg='Morsel should not be implicitly changed to dict on usage')
 
+    def test_traversal_filter(self):
+        data = [None, False, True, 0, 1, 0.0, 1.1, '', 'str', {}, {0: 0}, [], [1]]
+
+        self.assertEqual(
+            traverse_obj(data, (Ellipsis, filter)),
+            [True, 1, 1.1, 'str', {0: 0}, [1]],
+            '`filter` should filter falsy values')
+
     def test_get_first(self):
         self.assertEqual(get_first([{'a': None}, {'a': 'spam'}], 'a'), 'spam')
 

+ 6 - 0
youtube_dl/compat.py

@@ -3452,6 +3452,8 @@ except ImportError:
     except ImportError:
         compat_map = map
 
+
+# compat_filter, compat_filter_fns
 try:
     from future_builtins import filter as compat_filter
 except ImportError:
@@ -3459,6 +3461,9 @@ except ImportError:
         from itertools import ifilter as compat_filter
     except ImportError:
         compat_filter = filter
+# "Is this function one or maybe the other filter()?"
+compat_filter_fns = tuple(set((filter, compat_filter)))
+
 
 # compat_zip
 try:
@@ -3675,6 +3680,7 @@ __all__ = [
     'compat_etree_fromstring',
     'compat_etree_iterfind',
     'compat_filter',
+    'compat_filter_fns',
     'compat_get_terminal_size',
     'compat_getenv',
     'compat_getpass_getpass',

+ 8 - 0
youtube_dl/utils.py

@@ -53,6 +53,8 @@ from .compat import (
     compat_etree_fromstring,
     compat_etree_iterfind,
     compat_expanduser,
+    compat_filter as filter,
+    compat_filter_fns,
     compat_html_entities,
     compat_html_entities_html5,
     compat_http_client,
@@ -6283,6 +6285,7 @@ def traverse_obj(obj, *paths, **kwargs):
                             Read as: `{key: traverse_obj(obj, path) for key, path in dct.items()}`.
         - `any`-builtin:    Take the first matching object and return it, resetting branching.
         - `all`-builtin:    Take all matching objects and return them as a list, resetting branching.
+        - `filter`-builtin: Return the value if it is truthy, `None` otherwise.
 
         `tuple`, `list`, and `dict` all support nested paths and branches.
 
@@ -6497,6 +6500,11 @@ def traverse_obj(obj, *paths, **kwargs):
                     objs = (list(filtered_objs),)
                 continue
 
+            # filter might be from __builtin__, future_builtins, or itertools.ifilter
+            if key in compat_filter_fns:
+                objs = filter(None, objs)
+                continue
+
             if __debug__ and callable(key):
                 # Verify function signature
                 _try_bind_args(key, None, None)