tests: adjusted to updated probe module

tikuma-lsuhsc · tikuma-lsuhsc · commit 24bb4d4e87db · 2024-03-20T22:08:00.000-05:00
diff --git a/src/ffmpegio/probe.py b/src/ffmpegio/probe.py
@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 from typing import BinaryIO, Any, TypeAlias, Literal
+from numbers import Number
 from collections.abc import Sequence
 import json, re
 from fractions import Fraction
@@ -251,11 +252,11 @@ def full_details(
     show_streams: bool | None = True,
     show_programs: bool | None = False,
     show_chapters: bool | None = False,
-    select_streams: bool | None = None,
+    select_streams: str | int | None = None,
     keep_str_values: bool | None = False,
     cache_output: bool | None = False,
     sp_kwargs: dict[str, Any] | None = None,
-):
+) -> dict[str, str | Number | Fraction]:
     """Retrieve full details of a media file or stream
 
     :param url: URL of the media file/stream
@@ -268,8 +269,8 @@ def full_details(
     :type show_programs: bool, optional
     :param show_chapters: True to return chapter info, defaults to False
     :type show_chapters: bool, optional
-    :param select_streams: Indices of streams to get info of, defaults to None
-    :type select_streams: seq of int, optional
+    :param select_streams: Stream specifier of the streams to get info of, defaults to None to retrieve all
+    :type select_streams: str, int, optional
     :param keep_str_values: True to keep all field values as str,
                             defaults to False to convert numeric values
     :type keep_str_values: bool, optional
@@ -329,13 +330,25 @@ def format_basic(
     keep_str_values: bool | None = False,
     cache_output: bool | None = False,
     sp_kwargs: dict[str, Any] | None = None,
-):
+) -> dict[str, str | Number | Fraction]:
     """Retrieve basic media format info
 
     :param url: URL of the media file/stream
     :type url: str or seekable file-like object or bytes-like object
     :param entries: specify to narrow which information entries to retrieve. Default to None, to return all entries
     :type entries: seq of str
+    :param keep_optional_fields: True to return a missing optional field in the
+                        returned dict with None or "N/A" (if keep_str_values
+                        is True) as its value
+    :type keep_optional_fields: bool, optional
+    :param keep_str_values: True to keep all field values as str,
+                            defaults to False to convert numeric values
+    :type keep_str_values: bool, optional
+    :param cache_output: True to cache FFprobe output, defaults to False
+    :type cache_output: bool, optional
+    :param sp_kwargs: Additional keyword arguments for :py:func:`subprocess.run`,
+                      default to None
+    :type sp_kwargs: dict[str, Any], optional
     :return: set of media format information.
     :rtype: dict
 
@@ -380,13 +393,25 @@ def streams_basic(
     keep_str_values: bool | None = False,
     cache_output: bool | None = False,
     sp_kwargs: dict[str, Any] | None = None,
-):
+) -> dict[str, str | Number | Fraction]:
     """Retrieve basic info of media streams
 
     :param url: URL of the media file/stream
     :type url: str or seekable file-like object or bytes-like object
     :param entries: specify to narrow which stream entries to retrieve. Default to None, returning all entries
     :type entries: seq of str, optional
+    :param keep_optional_fields: True to return a missing optional field in the
+                        returned dict with None or "N/A" (if keep_str_values
+                        is True) as its value
+    :type keep_optional_fields: bool, optional
+    :param keep_str_values: True to keep all field values as str,
+                            defaults to False to convert numeric values
+    :type keep_str_values: bool, optional
+    :param cache_output: True to cache FFprobe output, defaults to False
+    :type cache_output: bool, optional
+    :param sp_kwargs: Additional keyword arguments for :py:func:`subprocess.run`,
+                      default to None
+    :type sp_kwargs: dict[str, Any], optional
     :return: List of media stream information.
     :rtype: list of dict
 
@@ -423,7 +448,7 @@ def video_streams_basic(
     keep_str_values: bool | None = False,
     cache_output: bool | None = False,
     sp_kwargs: dict[str, Any] | None = None,
-):
+) -> dict[str, str | Number | Fraction]:
     """Retrieve basic info of video streams
 
     :param url: URL of the media file/stream
@@ -432,6 +457,18 @@ def video_streams_basic(
     :type index: int, optional
     :param entries: specify to narrow which information entries to retrieve. Default to None, to return all entries
     :type entries: seq of str
+    :param keep_optional_fields: True to return a missing optional field in the
+                        returned dict with None or "N/A" (if keep_str_values
+                        is True) as its value
+    :type keep_optional_fields: bool, optional
+    :param keep_str_values: True to keep all field values as str,
+                            defaults to False to convert numeric values
+    :type keep_str_values: bool, optional
+    :param cache_output: True to cache FFprobe output, defaults to False
+    :type cache_output: bool, optional
+    :param sp_kwargs: Additional keyword arguments for :py:func:`subprocess.run`,
+                      default to None
+    :type sp_kwargs: dict[str, Any], optional
     :return: List of video stream information.
     :rtype: list of dict
 
@@ -528,7 +565,7 @@ def audio_streams_basic(
     keep_str_values: bool | None = False,
     cache_output: bool | None = False,
     sp_kwargs: dict[str, Any] | None = None,
-):
+) -> dict[str, str | Number | Fraction]:
     """Retrieve basic info of audio streams
 
     :param url: URL of the media file/stream
@@ -537,6 +574,18 @@ def audio_streams_basic(
     :type index: int, optional
     :param entries: specify to narrow which information entries to retrieve. Default to None, to return all entries
     :type entries: seq of str
+    :param keep_optional_fields: True to return a missing optional field in the
+                        returned dict with None or "N/A" (if keep_str_values
+                        is True) as its value
+    :type keep_optional_fields: bool, optional
+    :param keep_str_values: True to keep all field values as str,
+                            defaults to False to convert numeric values
+    :type keep_str_values: bool, optional
+    :param cache_output: True to cache FFprobe output, defaults to False
+    :type cache_output: bool, optional
+    :param sp_kwargs: Additional keyword arguments for :py:func:`subprocess.run`,
+                      default to None
+    :type sp_kwargs: dict[str, Any], optional
     :return: List of audio stream information.
     :rtype: list of dict
 
@@ -622,14 +671,14 @@ def query(
     keep_str_values: bool | None = False,
     cache_output: bool | None = False,
     sp_kwargs: dict[str, Any] | None = None,
-):
+) -> dict[str, Any] | Sequence[dict[str, Any]]:
     """Query specific fields of media format or stream
 
     :param url: URL of the media file/stream
     :type url: str or seekable file-like object or bytes-like object
     :param streams: stream specifier, defaults to None to get format
     :type streams: str, int, bool, optional
-    :param fields: info, defaults to None
+    :param fields: list of format/stream fields to retrieve, defaults to None (all fields)
     :type fields: sequence of str, optional
     :param keep_optional_fields: True to return a missing optional field in the
                         returned dict with None or "N/A" (if keep_str_values
@@ -755,24 +804,6 @@ def frames(
              if entries is str (i.e., a single entry)
     :rtype: list[dict] or list[str|int|float]
 
-    ``intervals`` argument
-    ----------------------
-
-    intervals argument can be specified in multiple ways to form the ``-read_intervals`` ffprobe option:
-
-    1) ``str`` - pass through the argument as-is to ffprobe
-    2) ``int`` - read this numbers of packets to read from the beginning of the file
-    3) ``float`` - read packets over this duration in seconds from the beginning of the file
-    4) ``seq[str|float, str|int|float]`` - sets start and end points
-       - start: str = as-is, float=starting time in seconds
-       - end: str = as-is, int=offset in # of packets, float=offset in seconds
-    5) ``dict`` - specifies start and end points with the following keys:
-       - 'start'        - (str|float) start time
-       - 'start_offset' - (str|float) start time offset from the previous read. Ignored if 'start' is present.
-       - 'end'          - (str|float) end time
-       - 'end_offset'   - (str|float|int) end time offset from the start time. Ignored if 'end' is present.
-    6) - ``seq[dict]``  - specify multiple intervals
-
     """
 
     is_single = isinstance(entries, str)
diff --git a/tests/test_audio.py b/tests/test_audio.py
@@ -71,10 +71,6 @@ def test_read_write():
     url = "tests/assets/testaudio-1m.mp3"
     outext = ".flac"
 
-    info = probe.audio_streams_basic(
-        url, index=0, entries=("sample_rate", "sample_fmt", "channels")
-    )[0]
-
     fs, x = audio.read(url)
 
     with tempfile.TemporaryDirectory() as tmpdirname:
diff --git a/tests/test_ffmpegprocess.py b/tests/test_ffmpegprocess.py
@@ -91,7 +91,9 @@ def test_popen():
 
 def test_popen_progress():
     url = "tests/assets/testvideo-1m.mp4"
-    info = probe.video_streams_basic(url, 0)[0]
+    info = probe.video_streams_basic(
+        url, 0, ["pix_fmt", "width", "height", "frame_rate"]
+    )
     pix_fmt_in = info["pix_fmt"]
     s_in = (info["width"], info["height"])
     r_in = info["frame_rate"]
diff --git a/tests/test_probe.py b/tests/test_probe.py
@@ -56,16 +56,16 @@ def test_query():
     )
 
     assert all(
-        st["bad_filed"] is None
+        st["max_bit_rate"] is None
         for st in probe.query(
-            url, "a", fields=("duration", "bad_filed"), return_none=True
+            url, "a", fields=("duration", "max_bit_rate"), keep_optional_fields=True
         )
     )
 
     assert (
-        probe.query(url, "v:0", fields=("duration", "bad_filed"), return_none=True)[
-            "bad_filed"
-        ]
+        probe.query(
+            url, "v:0", fields=("duration", "max_bit_rate"), keep_optional_fields=True
+        )["max_bit_rate"]
         is None
     )
 
@@ -79,7 +79,7 @@ def test_frames():
     info = probe.frames(
         url,
         streams="a:0",
-        intervals=10
+        intervals=10,
         # intervals='%+#20,30%+#15'
         # intervals=[{"end_offset": 20}, {"start": 30, "end_offset": 12}],
     )
@@ -89,7 +89,7 @@ def test_frames():
         url,
         "key_frame",
         "a:0",
-        intervals=10
+        intervals=10,
         # intervals='%+#20,30%+#15'
         # intervals=[{"end_offset": 20}, {"start": 30, "end_offset": 12}],
     )
@@ -99,7 +99,7 @@ def test_frames():
         "key_frame",
         "a:0",
         intervals=10,
-        accurate_time=True
+        accurate_time=True,
         # intervals='%+#20,30%+#15'
         # intervals=[{"end_offset": 20}, {"start": 30, "end_offset": 12}],
     )
@@ -110,7 +110,7 @@ def test_frames():
         url,
         streams="a:0",
         intervals=(23.4, 1),
-        accurate_time=True
+        accurate_time=True,
         # intervals='%+#20,30%+#15'
         # intervals=[{"end_offset": 20}, {"start": 30, "end_offset": 12}],
     )
