|
3 | 3 | from cython.cimports.av.stream import Stream |
4 | 4 |
|
5 | 5 |
|
6 | | -def _flatten(input_): |
7 | | - for x in input_: |
8 | | - if isinstance(x, (tuple, list)): |
9 | | - for y in _flatten(x): |
10 | | - yield y |
11 | | - else: |
12 | | - yield x |
13 | | - |
14 | | - |
15 | 6 | @cython.cfunc |
16 | 7 | def _get_media_type_enum(type: str) -> lib.AVMediaType: |
17 | 8 | if type == "video": |
@@ -131,35 +122,36 @@ def get(self, *args, **kwargs): |
131 | 122 | If nothing is selected, then all streams are returned. |
132 | 123 |
|
133 | 124 | """ |
| 125 | + selection: list = [] |
134 | 126 |
|
135 | | - selection = [] |
136 | | - |
137 | | - for x in _flatten((args, kwargs)): |
| 127 | + def process(x): |
138 | 128 | if x is None: |
139 | 129 | pass |
140 | | - |
141 | 130 | elif isinstance(x, Stream): |
142 | 131 | selection.append(x) |
143 | | - |
144 | 132 | elif isinstance(x, int): |
145 | 133 | selection.append(self._streams[x]) |
146 | | - |
| 134 | + elif isinstance(x, (tuple, list)): |
| 135 | + for item in x: |
| 136 | + process(item) |
147 | 137 | elif isinstance(x, dict): |
148 | 138 | for type_, indices in x.items(): |
149 | | - if ( |
150 | | - type_ == "streams" |
151 | | - ): # For compatibility with the pseudo signature |
152 | | - streams = self._streams |
153 | | - else: |
154 | | - streams = getattr(self, type_) |
| 139 | + # For compatibility with the pseudo signature |
| 140 | + streams = ( |
| 141 | + self._streams if type_ == "streams" else getattr(self, type_) |
| 142 | + ) |
155 | 143 | if not isinstance(indices, (tuple, list)): |
156 | 144 | indices = [indices] |
157 | 145 | for i in indices: |
158 | 146 | selection.append(streams[i]) |
159 | | - |
160 | 147 | else: |
161 | 148 | raise TypeError("Argument must be Stream or int.", type(x)) |
162 | 149 |
|
| 150 | + for arg in args: |
| 151 | + process(arg) |
| 152 | + if kwargs: |
| 153 | + process(kwargs) |
| 154 | + |
163 | 155 | return selection or self._streams[:] |
164 | 156 |
|
165 | 157 | def best(self, type: str, /, related: Stream | None = None): |
|
0 commit comments