From 21b5ec86c2c37d10c5bb97edd7051d3aac16bb3e Mon Sep 17 00:00:00 2001
From: Simon Sawicki <contact@grub4k.xyz>
Date: Mon, 24 Apr 2023 19:56:35 +0200
Subject: [PATCH] [utils] `traverse_obj`: Allow iterables in traversal (#6902)

Authored by: Grub4K
---
 test/test_utils.py | 4 ++++
 yt_dlp/utils.py    | 7 +++----
 2 files changed, 7 insertions(+), 4 deletions(-)

diff --git a/test/test_utils.py b/test/test_utils.py
index d4a301583..f2f3b8170 100644
--- a/test/test_utils.py
+++ b/test/test_utils.py
@@ -2016,6 +2016,8 @@ Line 1
                          msg='nested `...` queries should work')
         self.assertCountEqual(traverse_obj(_TEST_DATA, (..., ..., 'index')), range(4),
                               msg='`...` query result should be flattened')
+        self.assertEqual(traverse_obj(range(4), ...), list(range(4)),
+                         msg='`...` should accept iterables')
 
         # Test function as key
         self.assertEqual(traverse_obj(_TEST_DATA, lambda x, y: x == 'urls' and isinstance(y, list)),
@@ -2023,6 +2025,8 @@ Line 1
                          msg='function as query key should perform a filter based on (key, value)')
         self.assertCountEqual(traverse_obj(_TEST_DATA, lambda _, x: isinstance(x[0], str)), {'str'},
                               msg='exceptions in the query function should be catched')
+        self.assertEqual(traverse_obj(range(4), lambda _, x: x % 2 == 0), [0, 2],
+                         msg='function key should accept iterables')
         if __debug__:
             with self.assertRaises(Exception, msg='Wrong function signature should raise in debug'):
                 traverse_obj(_TEST_DATA, lambda a: ...)
diff --git a/yt_dlp/utils.py b/yt_dlp/utils.py
index 746a2885d..f69311462 100644
--- a/yt_dlp/utils.py
+++ b/yt_dlp/utils.py
@@ -5528,7 +5528,6 @@ def traverse_obj(
                             If no `default` is given and the last path branches, a `list` of results
                             is always returned. If a path ends on a `dict` that result will always be a `dict`.
     """
-    is_sequence = lambda x: isinstance(x, collections.abc.Sequence) and not isinstance(x, (str, bytes))
     casefold = lambda k: k.casefold() if isinstance(k, str) else k
 
     if isinstance(expected_type, type):
@@ -5564,7 +5563,7 @@ def traverse_obj(
             branching = True
             if isinstance(obj, collections.abc.Mapping):
                 result = obj.values()
-            elif is_sequence(obj):
+            elif isinstance(obj, collections.abc.Iterable) and not isinstance(obj, (str, bytes)):
                 result = obj
             elif isinstance(obj, re.Match):
                 result = obj.groups()
@@ -5578,7 +5577,7 @@ def traverse_obj(
             branching = True
             if isinstance(obj, collections.abc.Mapping):
                 iter_obj = obj.items()
-            elif is_sequence(obj):
+            elif isinstance(obj, collections.abc.Iterable) and not isinstance(obj, (str, bytes)):
                 iter_obj = enumerate(obj)
             elif isinstance(obj, re.Match):
                 iter_obj = itertools.chain(
@@ -5614,7 +5613,7 @@ def traverse_obj(
                 result = next((v for k, v in obj.groupdict().items() if casefold(k) == key), None)
 
         elif isinstance(key, (int, slice)):
-            if is_sequence(obj):
+            if isinstance(obj, collections.abc.Sequence) and not isinstance(obj, (str, bytes)):
                 branching = isinstance(key, slice)
                 with contextlib.suppress(IndexError):
                     result = obj[key]