diff --git a/av/audio/stream.pyi b/av/audio/stream.pyi index 00f4fd04f..e14323f2a 100644 --- a/av/audio/stream.pyi +++ b/av/audio/stream.pyi @@ -8,9 +8,9 @@ from .format import AudioFormat from .frame import AudioFrame class AudioStream(Stream): + type: Literal["audio"] format: AudioFormat codec_context: AudioCodecContext - type: Literal["audio"] def encode(self, frame: AudioFrame | None = None) -> list[Packet]: ... def decode(self, packet: Packet | None = None) -> list[AudioFrame]: ... diff --git a/av/video/format.pyi b/av/video/format.pyi index 33aa8d3c2..16739a0e3 100644 --- a/av/video/format.pyi +++ b/av/video/format.pyi @@ -8,7 +8,7 @@ class VideoFormat: is_planar: bool is_rgb: bool - def __init__(self, name: str, width: int = 0, height: int = 0): ... + def __init__(self, name: str, width: int = 0, height: int = 0) -> None: ... def chroma_width(self, luma_width: int = 0) -> int: ... def chroma_height(self, luma_height: int = 0) -> int: ... @@ -21,4 +21,4 @@ class VideoFormatComponent: width: int height: int - def __init__(self, format: VideoFormat, index: int): ... + def __init__(self, format: VideoFormat, index: int) -> None: ... diff --git a/av/video/frame.pyi b/av/video/frame.pyi index a387ad37f..da6575328 100644 --- a/av/video/frame.pyi +++ b/av/video/frame.pyi @@ -1,3 +1,5 @@ +from typing import Any, Union + import numpy as np from PIL import Image @@ -6,6 +8,13 @@ from av.frame import Frame from .format import VideoFormat from .plane import VideoPlane +from .reformatter import ColorRange, Colorspace + +_SupportedNDarray = Union[ + np.ndarray[Any, np.dtype[np.uint8]], + np.ndarray[Any, np.dtype[np.uint16]], + np.ndarray[Any, np.dtype[np.float32]], +] class PictureType(EnumItem): NONE: int @@ -27,26 +36,31 @@ class VideoFrame(Frame): key_frame: bool interlaced_frame: bool pict_type: int + colorspace: int + color_range: int - @staticmethod - def from_image(img: Image.Image) -> VideoFrame: ... - @staticmethod - def from_ndarray(array: np.ndarray, format: str = "rgb24") -> VideoFrame: ... - @staticmethod - def from_numpy_buffer(array: np.ndarray, format: str = "rgb24"): ... def __init__( - self, name: str, width: int = 0, height: int = 0, format: str = "yuv420p" - ): ... - def to_image(self, **kwargs) -> Image.Image: ... - def to_ndarray(self, **kwargs) -> np.ndarray: ... + self, width: int = 0, height: int = 0, format: str = "yuv420p" + ) -> None: ... def reformat( self, width: int | None = None, height: int | None = None, format: str | None = None, - src_colorspace=None, - dst_colorspace=None, + src_colorspace: Colorspace | None = None, + dst_colorspace: Colorspace | None = None, interpolation: int | str | None = None, src_color_range: int | str | None = None, dst_color_range: int | str | None = None, ) -> VideoFrame: ... + def to_rgb(self, **kwargs: Any) -> VideoFrame: ... + def to_image(self, **kwargs: Any) -> Image.Image: ... + def to_ndarray(self, **kwargs: Any) -> _SupportedNDarray: ... + @staticmethod + def from_image(img: Image.Image) -> VideoFrame: ... + @staticmethod + def from_numpy_buffer( + array: _SupportedNDarray, format: str = "rgb24" + ) -> VideoFrame: ... + @staticmethod + def from_ndarray(array: _SupportedNDarray, format: str = "rgb24") -> VideoFrame: ... diff --git a/av/video/plane.pyi b/av/video/plane.pyi index 592d3eb15..9bc1f0d77 100644 --- a/av/video/plane.pyi +++ b/av/video/plane.pyi @@ -6,4 +6,4 @@ class VideoPlane: height: int buffer_size: int - def __init__(self, frame: VideoFrame, index: int): ... + def __init__(self, frame: VideoFrame, index: int) -> None: ... diff --git a/av/video/reformatter.pyi b/av/video/reformatter.pyi index f17055a4f..a68dc9718 100644 --- a/av/video/reformatter.pyi +++ b/av/video/reformatter.pyi @@ -1,6 +1,5 @@ from av.enum import EnumItem -from .format import VideoFormat from .frame import VideoFrame class Interpolation(EnumItem): @@ -44,8 +43,8 @@ class VideoReformatter: width: int | None = None, height: int | None = None, format: str | None = None, - src_colorspace=None, - dst_colorspace=None, + src_colorspace: Colorspace | None = None, + dst_colorspace: Colorspace | None = None, interpolation: int | str | None = None, src_color_range: int | str | None = None, dst_color_range: int | str | None = None, diff --git a/av/video/stream.pyi b/av/video/stream.pyi index f6be52460..2ef4adf89 100644 --- a/av/video/stream.pyi +++ b/av/video/stream.pyi @@ -1,19 +1,21 @@ from fractions import Fraction -from typing import Any +from typing import Any, Literal from av.packet import Packet from av.stream import Stream from .codeccontext import VideoCodecContext +from .format import VideoFormat from .frame import VideoFrame class VideoStream(Stream): - name: str width: int height: int + format: VideoFormat pix_fmt: str | None sample_aspect_ratio: Fraction | None codec_context: VideoCodecContext + type: Literal["video"] # from codec context bit_rate: int | None