Skip to content

Commit db710df

Browse files
committed
Simplify StreamContainer.get method
1 parent 26a489c commit db710df

File tree

1 file changed

+14
-22
lines changed

1 file changed

+14
-22
lines changed

av/container/streams.py

Lines changed: 14 additions & 22 deletions
Original file line numberDiff line numberDiff line change
@@ -3,15 +3,6 @@
33
from cython.cimports.av.stream import Stream
44

55

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-
156
@cython.cfunc
167
def _get_media_type_enum(type: str) -> lib.AVMediaType:
178
if type == "video":
@@ -131,35 +122,36 @@ def get(self, *args, **kwargs):
131122
If nothing is selected, then all streams are returned.
132123
133124
"""
125+
selection: list = []
134126

135-
selection = []
136-
137-
for x in _flatten((args, kwargs)):
127+
def process(x):
138128
if x is None:
139129
pass
140-
141130
elif isinstance(x, Stream):
142131
selection.append(x)
143-
144132
elif isinstance(x, int):
145133
selection.append(self._streams[x])
146-
134+
elif isinstance(x, (tuple, list)):
135+
for item in x:
136+
process(item)
147137
elif isinstance(x, dict):
148138
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+
)
155143
if not isinstance(indices, (tuple, list)):
156144
indices = [indices]
157145
for i in indices:
158146
selection.append(streams[i])
159-
160147
else:
161148
raise TypeError("Argument must be Stream or int.", type(x))
162149

150+
for arg in args:
151+
process(arg)
152+
if kwargs:
153+
process(kwargs)
154+
163155
return selection or self._streams[:]
164156

165157
def best(self, type: str, /, related: Stream | None = None):

0 commit comments

Comments
 (0)