Continue test refactor

Author: WyattBlue

Makefile

@@ -30,7 +30,7 @@ fate-suite:
 lint:
 	$(PIP) install -U black isort flake8 flake8-pyproject pillow numpy mypy==1.11.2
 	black --check av examples tests setup.py
-	flake8 av examples tests
+	flake8 av
 	isort --check-only --diff av examples tests
 	mypy av tests
 

av/__init__.py

@@ -15,33 +15,62 @@
         + os.environ["PATH"]
     )
 
-# MUST import the core before anything else in order to initalize the underlying
+# MUST import the core before anything else in order to initialize the underlying
 # library that is being wrapped.
 from av._core import time_base, library_versions
 
 # Capture logging (by importing it).
 from av import logging
 
-# For convenience, IMPORT ALL OF THE THINGS (that are constructable by the user).
+# For convenience, import all common attributes.
 from av.about import __version__
+from av.audio.codeccontext import AudioCodecContext
 from av.audio.fifo import AudioFifo
 from av.audio.format import AudioFormat
 from av.audio.frame import AudioFrame
 from av.audio.layout import AudioLayout
 from av.audio.resampler import AudioResampler
+from av.audio.stream import AudioStream
 from av.bitstream import BitStreamFilterContext, bitstream_filters_available
 from av.codec.codec import Codec, codecs_available
 from av.codec.context import CodecContext
 from av.container import open
 from av.format import ContainerFormat, formats_available
 from av.packet import Packet
 from av.error import *  # noqa: F403; This is limited to exception types.
+from av.video.codeccontext import VideoCodecContext
 from av.video.format import VideoFormat
 from av.video.frame import VideoFrame
+from av.video.stream import VideoStream
 
 # Backwards compatibility
 AVError = FFmpegError  # noqa: F405
 
+__all__ = (
+    "time_base",
+    "library_versions",
+    "AudioCodecContext",
+    "AudioFifo",
+    "AudioFormat",
+    "AudioFrame",
+    "AudioLayout",
+    "AudioResampler",
+    "AudioStream",
+    "BitStreamFilterContext",
+    "bitstream_filters_available",
+    "Codec",
+    "codecs_available",
+    "CodecContext",
+    "open",
+    "ContainerFormat",
+    "formats_available",
+    "Packet",
+    "VideoCodecContext",
+    "VideoFormat",
+    "VideoFrame",
+    "VideoStream",
+)
+
 
 def get_include() -> str:
     """

av/container/core.pyi

@@ -1,4 +1,4 @@
-from numbers import Real
+from fractions import Fraction
 from pathlib import Path
 from types import TracebackType
 from typing import Any, Callable, Literal, Type, overload
@@ -10,6 +10,8 @@ from .input import InputContainer
 from .output import OutputContainer
 from .streams import StreamContainer
 
+Real = int | float | Fraction
+
 class Flags(EnumFlag):
     GENPTS: int
     IGNIDX: int

tests/common.py

@@ -14,6 +14,8 @@
 except ImportError:
     has_pillow = False
 
+__all__ = ("fate_suite",)
+
 
 is_windows = os.name == "nt"
 skip_tests = frozenset(os.environ.get("PYAV_SKIP_TESTS", "").split(","))

tests/test_codec_context.py

@@ -2,17 +2,48 @@
 
 import os
 from fractions import Fraction
+from typing import Iterator, TypedDict, overload
 from unittest import SkipTest
 
 import av
-from av import AudioLayout, AudioResampler, Codec, Packet
+from av import (
+    AudioCodecContext,
+    AudioFrame,
+    AudioLayout,
+    AudioResampler,
+    Codec,
+    Packet,
+    VideoCodecContext,
+    VideoFrame,
+)
 from av.codec.codec import UnknownCodecError
 from av.video.frame import PictureType
 
 from .common import TestCase, fate_suite
 
 
-def iter_raw_frames(path, packet_sizes, ctx):
+class Options(TypedDict, total=False):
+    b: str
+    crf: str
+    pix_fmt: str
+    width: int
+    height: int
+    max_frames: int
+    time_base: Fraction
+    gop_size: int
+
+
+@overload
+def iter_raw_frames(
+    path: str, packet_sizes: list, ctx: VideoCodecContext
+) -> Iterator[VideoFrame]: ...
+@overload
+def iter_raw_frames(
+    path: str, packet_sizes: list, ctx: AudioCodecContext
+) -> Iterator[AudioFrame]: ...
+def iter_raw_frames(
+    path: str, packet_sizes: list, ctx: VideoCodecContext | AudioCodecContext
+) -> Iterator[VideoFrame | AudioFrame]:
     with open(path, "rb") as f:
         for i, size in enumerate(packet_sizes):
             packet = Packet(size)
@@ -244,24 +275,24 @@ def image_sequence_encode(self, codec_name: str) -> None:
                 assert frame.height == height
                 assert frame.format.name == pix_fmt
 
-    def test_encoding_h264(self):
+    def test_encoding_h264(self) -> None:
         self.video_encoding("h264", {"crf": "19"})
 
-    def test_encoding_mpeg4(self):
+    def test_encoding_mpeg4(self) -> None:
         self.video_encoding("mpeg4")
 
-    def test_encoding_xvid(self):
+    def test_encoding_xvid(self) -> None:
         self.video_encoding("mpeg4", codec_tag="xvid")
 
-    def test_encoding_mpeg1video(self):
+    def test_encoding_mpeg1video(self) -> None:
         self.video_encoding("mpeg1video")
 
-    def test_encoding_dvvideo(self):
-        options = {"pix_fmt": "yuv411p", "width": 720, "height": 480}
+    def test_encoding_dvvideo(self) -> None:
+        options: Options = {"pix_fmt": "yuv411p", "width": 720, "height": 480}
         self.video_encoding("dvvideo", options)
 
-    def test_encoding_dnxhd(self):
-        options = {
+    def test_encoding_dnxhd(self) -> None:
+        options: Options = {
             "b": "90M",  # bitrate
             "pix_fmt": "yuv422p",
             "width": 1920,
@@ -271,7 +302,12 @@ def test_encoding_dnxhd(self):
         }
         self.video_encoding("dnxhd", options)
 
-    def video_encoding(self, codec_name, options={}, codec_tag=None):
+    def video_encoding(
+        self,
+        codec_name: str,
+        options: Options = {},
+        codec_tag: str | None = None,
+    ) -> None:
         try:
             codec = Codec(codec_name, "w")
         except UnknownCodecError:
@@ -280,21 +316,23 @@ def video_encoding(self, codec_name, options={}, codec_tag=None):
         container = av.open(fate_suite("h264/interlaced_crop.mp4"))
         video_stream = container.streams.video[0]
 
+        assert video_stream.time_base is not None
+
         pix_fmt = options.pop("pix_fmt", "yuv420p")
         width = options.pop("width", 640)
         height = options.pop("height", 480)
         max_frames = options.pop("max_frames", 50)
         time_base = options.pop("time_base", video_stream.time_base)
         gop_size = options.pop("gop_size", 20)
 
-        ctx = codec.create()
+        ctx = codec.create("video")
         ctx.width = width
         ctx.height = height
         ctx.time_base = time_base
         ctx.framerate = 1 / ctx.time_base
         ctx.pix_fmt = pix_fmt
         ctx.gop_size = gop_size
-        ctx.options = options  # TODO
+        ctx.options = options  # type: ignore
         if codec_tag:
             ctx.codec_tag = codec_tag
         ctx.open()
@@ -326,7 +364,7 @@ def video_encoding(self, codec_name, options={}, codec_tag=None):
         if codec_name == "libx264":
             dec_codec_name = "h264"
 
-        ctx = av.Codec(dec_codec_name, "r").create()
+        ctx = av.Codec(dec_codec_name, "r").create("video")
         ctx.open()
 
         keyframe_indices = []
@@ -341,7 +379,7 @@ def video_encoding(self, codec_name, options={}, codec_tag=None):
 
         assert frame_count == decoded_frame_count
 
-        self.assertIsInstance(
+        assert isinstance(
             all(keyframe_index for keyframe_index in keyframe_indices), int
         )
         decoded_gop_sizes = [
@@ -350,7 +388,7 @@ def video_encoding(self, codec_name, options={}, codec_tag=None):
         if codec_name in ("dvvideo", "dnxhd") and all(
             i == 1 for i in decoded_gop_sizes
         ):
-            raise SkipTest()
+            raise SkipTest("I'm not sure why we skip this actually.")
         for i in decoded_gop_sizes:
             assert i == gop_size
 

tests/test_decode.py

@@ -51,8 +51,10 @@ def test_decode_audio_sample_count(self) -> None:
         assert audio_stream.duration is not None
         assert audio_stream.time_base is not None
         total_samples = (
-            audio_stream.duration * audio_stream.sample_rate.numerator
-        ) / audio_stream.time_base.denominator
+            audio_stream.duration
+            * audio_stream.sample_rate.numerator
+            / audio_stream.time_base.denominator
+        )
         assert sample_count == total_samples
 
     def test_decoded_time_base(self):
@@ -125,11 +127,9 @@ def test_decode_close_then_use(self):
                 except AssertionError:
                     pass
 
-    def test_flush_decoded_video_frame_count(self):
+    def test_flush_decoded_video_frame_count(self) -> None:
         container = av.open(fate_suite("h264/interlaced_crop.mp4"))
-        video_stream = next(s for s in container.streams if s.type == "video")
-
-        self.assertIs(video_stream, container.streams.video[0])
+        video_stream = container.streams.video[0]
 
         # Decode the first GOP, which requires a flush to get all frames
         have_keyframe = False
@@ -148,11 +148,11 @@ def test_flush_decoded_video_frame_count(self):
                 output_count += 1
 
         # Check the test works as expected and requires a flush
-        self.assertLess(output_count, input_count)
+        assert output_count < input_count
 
         for frame in video_stream.decode(None):
             # The Frame._time_base is not set by PyAV
-            self.assertIsNone(frame.time_base)
+            assert frame.time_base is None
             output_count += 1
 
         assert output_count == input_count

tests/test_enums.py

@@ -15,15 +15,13 @@ def define_foobar(self, **kwargs):
     def test_basics(self):
         cls = self.define_foobar()
 
-        self.assertIsInstance(cls, EnumType)
+        assert isinstance(cls, EnumType)
 
         foo = cls.FOO
 
-        self.assertIsInstance(foo, cls)
-        self.assertEqual(foo.name, "FOO")
-        self.assertEqual(foo.value, 1)
-
-        self.assertNotIsInstance(foo, PickleableFooBar)
+        assert isinstance(foo, cls)
+        assert foo.name == "FOO" and foo.value == 1
+        assert not isinstance(foo, PickleableFooBar)
 
     def test_access(self):
         cls = self.define_foobar()
@@ -35,9 +33,9 @@ def test_access(self):
         assert foo1 is foo3
         assert foo1 is foo4
 
-        self.assertIn(foo1, cls)
-        self.assertIn("FOO", cls)
-        self.assertIn(1, cls)
+        assert foo1 in cls
+        assert "FOO" in cls
+        assert 1 in cls
 
         self.assertRaises(KeyError, lambda: cls["not a foo"])
         self.assertRaises(KeyError, lambda: cls[10])
@@ -53,12 +51,12 @@ def test_casting(self):
         self.assertEqual(repr(foo), "<tests.test_enums.Foobar:FOO(0x1)>")
 
         str_foo = str(foo)
-        self.assertIsInstance(str_foo, str)
-        self.assertEqual(str_foo, "FOO")
+        assert isinstance(str_foo, str)
+        assert str_foo == "FOO"
 
         int_foo = int(foo)
-        self.assertIsInstance(int_foo, int)
-        self.assertEqual(int_foo, 1)
+        assert isinstance(int_foo, int)
+        assert int_foo == 1
 
     def test_iteration(self):
         cls = self.define_foobar()