Repository: yt-dlp/yt-dlp
Pull request #5174 (2022-10-07 23:36:16)

[utils] `traverse_obj`: Allow `re.Match` objects
**IMPORTANT**: PRs without the template will be CLOSED

### Description of your *pull request* and other information

</details>

This PR implements traversal into `re.Match` objects and related tests.
These keys are supported:
- `...`: returns `match.groups()`
- `function`: gives groupno, value as parameters (group 0 is the whole match).
- `str`/`int`: returns `match.group(key)`. Supports `casesense`.

<details open><summary>Template</summary> <!-- OPEN is intentional -->

<!--

# PLEASE FOLLOW THE GUIDE BELOW

- You will be asked some questions, please read them **carefully** and answer honestly
- Put an `x` into all the boxes `[ ]` relevant to your *pull request* (like [x])
- Use *Preview* tab to see how your *pull request* will actually look like

-->

### Before submitting a *pull request* make sure you have:
- [x] At least skimmed through [contributing guidelines](https://github.com/yt-dlp/yt-dlp/blob/master/CONTRIBUTING.md#developer-instructions) including [yt-dlp coding conventions](https://github.com/yt-dlp/yt-dlp/blob/master/CONTRIBUTING.md#yt-dlp-coding-conventions)
- [x] [Searched](https://github.com/yt-dlp/yt-dlp/search?q=is%3Apr&type=Issues) the bugtracker for similar pull requests
- [x] Checked the code with [flake8](https://pypi.python.org/pypi/flake8) and [ran relevant tests](https://github.com/yt-dlp/yt-dlp/blob/master/CONTRIBUTING.md#developer-instructions)

### In order to be accepted and merged into yt-dlp each piece of code must be in public domain or released under [Unlicense](http://unlicense.org/). Check one of the following options:
- [x] I am the original author of this code and I am willing to release it under [Unlicense](http://unlicense.org/)
- [ ] I am not the original author of this code but it is in public domain or released under [Unlicense](http://unlicense.org/) (provide reliable evidence)

### What is the purpose of your *pull request*?
- [ ] Fix or improvement to an extractor (Make sure to add/update tests)
- [ ] New extractor ([Piracy websites will not be accepted](https://github.com/yt-dlp/yt-dlp/blob/master/CONTRIBUTING.md#is-the-website-primarily-used-for-piracy))
- [x] Core bug fix/improvement
- [ ] New feature (It is strongly [recommended to open an issue first](https://github.com/yt-dlp/yt-dlp/blob/master/CONTRIBUTING.md#adding-new-feature-or-making-overarching-changes))


pukkandan: You missed docs
Grub4K: Right, but also we have not noted what input objects we supported before as well...
I would add all supported objects into the docs as well as `re.Match`.
pukkandan: https://github.com/yt-dlp/yt-dlp/blob/776bb387719c474a8f7a87e6e01799fc4f8859e3/yt_dlp/utils.py#L5313

so far, this acts as a good-enough documentation

But sure, documenting it explicitly sounds good

Patch:
diff --git a/yt_dlp/utils.py b/yt_dlp/utils.py
index 7d8e97162..cb14908c7 100644
--- a/yt_dlp/utils.py
+++ b/yt_dlp/utils.py
@@ -5305,13 +5305,14 @@ def traverse_obj(
 
     Each of the provided `paths` is tested and the first producing a valid result will be returned.
     The next path will also be tested if the path branched but no results could be found.
+    Supported values for traversal are `Mapping`, `Sequence` and `re.Match`.
     A value of None is treated as the absence of a value.
 
     The paths will be wrapped in `variadic`, so that `'key'` is conveniently the same as `('key', )`.
 
     The keys in the path can be one of:
         - `None`:           Return the current object.
-        - `str`/`int`:      Return `obj[key]`.
+        - `str`/`int`:      Return `obj[key]`. For `re.Match, return `obj.group(key)`.
         - `slice`:          Branch out and return all values in `obj[key]`.
         - `Ellipsis`:       Branch out and return a list of all values.
         - `tuple`/`list`:   Branch out and return a list of all matching values.
@@ -5322,7 +5323,7 @@ def traverse_obj(
         - `dict`            Transform the current object and return a matching dict.
                             Read as: `{key: traverse_obj(obj, path) for key, path in dct.items()}`.
 
-        `tuple`, `list`, and `dict` all support nested paths and branches
+        `tuple`, `list`, and `dict` all support nested paths and branches.
 
     @params paths           Paths which to traverse by.
     @param default          Value to return if the paths do not match.
@@ -5370,6 +5371,8 @@ def apply_key(key, obj):
                 yield from obj.values()
             elif is_sequence(obj):
                 yield from obj
+            elif isinstance(obj, re.Match):
+                yield from obj.groups()
             elif traverse_string:
                 yield from str(obj)
 
@@ -5378,6 +5381,8 @@ def apply_key(key, obj):
                 iter_obj = enumerate(obj)
             elif isinstance(obj, collections.abc.Mapping):
                 iter_obj = obj.items()
+            elif isinstance(obj, re.Match):
+                iter_obj = enumerate((obj.group(), *obj.groups()))
             elif traverse_string:
                 iter_obj = enumerate(str(obj))
             else:
@@ -5389,10 +5394,21 @@ def apply_key(key, obj):
             yield {k: v if v is not None else default for k, v in iter_obj
                    if v is not None or default is not NO_DEFAULT}
 
-        elif isinstance(obj, dict):
+        elif isinstance(obj, collections.abc.Mapping):
             yield (obj.get(key) if casesense or (key in obj)
                    else next((v for k, v in obj.items() if casefold(k) == key), None))
 
+        elif isinstance(obj, re.Match):
+            if isinstance(key, int) or casesense:
+                with contextlib.suppress(IndexError):
+                    yield obj.group(key)
+                    return
+
+            if not isinstance(key, str):
+                return
+
+            yield next((v for k, v in obj.groupdict().items() if casefold(k) == key), None)
+
         else:
             if is_user_input:
                 key = (int_or_none(key) if ':' not in key


Test patch:
diff --git a/test/test_utils.py b/test/test_utils.py
index 6f3f6cb91..90085a9c0 100644
--- a/test/test_utils.py
+++ b/test/test_utils.py
@@ -2,6 +2,7 @@
 
 # Allow direct execution
 import os
+import re
 import sys
 import unittest
 
@@ -2080,6 +2081,25 @@ def test_traverse_obj(self):
         with self.assertRaises(TypeError, msg='too many params should result in error'):
             traverse_obj(_IS_USER_INPUT_DATA, ('range8', ':::'), is_user_input=True)
 
+        # Test re.Match as input obj
+        mobj = re.fullmatch(r'0(12)(?P<group>3)(4)?', '0123')
+        self.assertEqual(traverse_obj(mobj, ...), [x for x in mobj.groups() if x is not None],
+                         msg='`...` on a `re.Match` should give its `groups()`')
+        self.assertEqual(traverse_obj(mobj, lambda k, _: k in (0, 2)), ['0123', '3'],
+                         msg='function on a `re.Match` should give groupno, value starting at 0')
+        self.assertEqual(traverse_obj(mobj, 'group'), '3',
+                         msg='str key on a `re.Match` should give group with that name')
+        self.assertEqual(traverse_obj(mobj, 2), '3',
+                         msg='int key on a `re.Match` should give group with that name')
+        self.assertEqual(traverse_obj(mobj, 'gRoUp', casesense=False), '3',
+                         msg='str key on a `re.Match` should respect casesense')
+        self.assertEqual(traverse_obj(mobj, 'fail'), None,
+                         msg='failing str key on a `re.Match` should return `default`')
+        self.assertEqual(traverse_obj(mobj, 'gRoUpS', casesense=False), None,
+                         msg='failing str key on a `re.Match` should return `default`')
+        self.assertEqual(traverse_obj(mobj, 8), None,
+                         msg='failing int key on a `re.Match` should return `default`')
+
 
 if __name__ == '__main__':
     unittest.main()