From 2a212dff25f2a20b771d32f066188fffc90210eb Mon Sep 17 00:00:00 2001 From: WyattBlue Date: Mon, 19 Jan 2026 03:58:39 -0500 Subject: [PATCH] Simplify StreamContainer.get method --- av/container/streams.py | 36 ++++++++++++++---------------------- 1 file changed, 14 insertions(+), 22 deletions(-) diff --git a/av/container/streams.py b/av/container/streams.py index 027c0016b..0045b5c85 100644 --- a/av/container/streams.py +++ b/av/container/streams.py @@ -3,15 +3,6 @@ from cython.cimports.av.stream import Stream -def _flatten(input_): - for x in input_: - if isinstance(x, (tuple, list)): - for y in _flatten(x): - yield y - else: - yield x - - @cython.cfunc def _get_media_type_enum(type: str) -> lib.AVMediaType: if type == "video": @@ -131,35 +122,36 @@ def get(self, *args, **kwargs): If nothing is selected, then all streams are returned. """ + selection: list = [] - selection = [] - - for x in _flatten((args, kwargs)): + def process(x): if x is None: pass - elif isinstance(x, Stream): selection.append(x) - elif isinstance(x, int): selection.append(self._streams[x]) - + elif isinstance(x, (tuple, list)): + for item in x: + process(item) elif isinstance(x, dict): for type_, indices in x.items(): - if ( - type_ == "streams" - ): # For compatibility with the pseudo signature - streams = self._streams - else: - streams = getattr(self, type_) + # For compatibility with the pseudo signature + streams = ( + self._streams if type_ == "streams" else getattr(self, type_) + ) if not isinstance(indices, (tuple, list)): indices = [indices] for i in indices: selection.append(streams[i]) - else: raise TypeError("Argument must be Stream or int.", type(x)) + for arg in args: + process(arg) + if kwargs: + process(kwargs) + return selection or self._streams[:] def best(self, type: str, /, related: Stream | None = None):