Add stubs for D3DShot (#8652)

Author: Avasam

stubs/D3DShot/METADATA.toml

@@ -0,0 +1,7 @@
+version = "0.1.*"
+requires = ["types-Pillow"]
+
+[tool.stubtest]
+# The library only works on Windows; we currently only run stubtest on Ubuntu for third-party stubs in CI.
+# See #8660
+skip = true

stubs/D3DShot/d3dshot/__init__.pyi

@@ -0,0 +1,12 @@
+from d3dshot.capture_output import CaptureOutputs as CaptureOutputs
+from d3dshot.d3dshot import D3DShot as D3DShot
+
+pil_is_available: bool
+numpy_is_available: bool
+pytorch_is_available: bool
+pytorch_gpu_is_available: bool
+capture_output_mapping: dict[str, CaptureOutputs]
+capture_outputs: list[str]
+
+def determine_available_capture_outputs() -> list[CaptureOutputs]: ...
+def create(capture_output: str = ..., frame_buffer_size: int = ...) -> D3DShot: ...

stubs/D3DShot/d3dshot/capture_output.pyi

@@ -0,0 +1,44 @@
+import enum
+from _typeshed import Incomplete
+from collections.abc import Sequence
+from ctypes import _CVoidConstPLike
+from typing_extensions import Literal, TypeAlias
+
+from PIL import Image
+
+_Frame: TypeAlias = Image.Image | Incomplete
+# TODO: Complete types once we can import non-types dependencies
+# See: #5768
+# from torch import Tensor
+# from comtypes import IUnknown
+# import numpy.typing as npt
+# _Frame: TypeAlias = Image.Image | npt.NDArray[np.int32] | npt.NDArray[np.float32] | Tensor
+
+class CaptureOutputs(enum.Enum):
+    PIL: int
+    NUMPY: int
+    NUMPY_FLOAT: int
+    PYTORCH: int
+    PYTORCH_FLOAT: int
+    PYTORCH_GPU: int
+    PYTORCH_FLOAT_GPU: int
+
+class CaptureOutputError(BaseException): ...
+
+# All CaptureOutput methods just reference the backend. Making this both a base class and a wrapper.
+class CaptureOutput:
+    # `backend` is a subclass of CaptureOutput based on the CaptureOutputs enum passed to __init__
+    backend: CaptureOutput
+    def __init__(self, backend: CaptureOutputs = ...) -> None: ...
+    def process(
+        self,
+        pointer: _CVoidConstPLike,
+        pitch: int,
+        size: int,
+        width: int,
+        height: int,
+        region: tuple[int, int, int, int],
+        rotation: int,
+    ) -> _Frame: ...
+    def to_pil(self, frame: _Frame) -> Image.Image: ...
+    def stack(self, frames: Sequence[_Frame], stack_dimension: Literal["first", "last"]) -> _Frame: ...

stubs/D3DShot/d3dshot/capture_outputs/__init__.pyi

@@ -0,0 +1,29 @@
+from _typeshed import Incomplete
+from collections.abc import Sequence
+from ctypes import _CVoidConstPLike
+from typing_extensions import Literal, TypeAlias
+
+from d3dshot.capture_output import CaptureOutput
+from PIL import Image
+
+# TODO: Complete types once we can import non-types dependencies
+# See: #5768
+# import numpy as np
+# import numpy.typing as npt
+# _NDArray: TypeAlias = npt.NDArray[np.int32]
+_NDArray: TypeAlias = Incomplete
+
+class NumpyCaptureOutput(CaptureOutput):
+    def __init__(self) -> None: ...
+    def process(
+        self,
+        pointer: _CVoidConstPLike,
+        pitch: int,
+        size: int,
+        width: int,
+        height: int,
+        region: tuple[int, int, int, int],
+        rotation: int,
+    ) -> _NDArray: ...
+    def to_pil(self, frame: _NDArray) -> Image.Image: ...
+    def stack(self, frames: Sequence[_NDArray] | _NDArray, stack_dimension: Literal["first", "last"]) -> _NDArray: ...

stubs/D3DShot/d3dshot/capture_outputs/numpy_capture_output.pyi

@@ -0,0 +1,5 @@
+from d3dshot.capture_outputs.numpy_capture_output import NumpyCaptureOutput
+
+# TODO: Once we can import non-types dependencies, this CaptureOutput should be float based
+# See: #5768
+class NumpyFloatCaptureOutput(NumpyCaptureOutput): ...

stubs/D3DShot/d3dshot/capture_outputs/numpy_float_capture_output.pyi

@@ -0,0 +1,25 @@
+from collections.abc import Sequence
+from ctypes import _CVoidConstPLike
+from typing import TypeVar
+from typing_extensions import TypeAlias
+
+from d3dshot.capture_output import CaptureOutput
+from PIL import Image
+
+_Unused: TypeAlias = object
+_ImageT = TypeVar("_ImageT", bound=Image.Image)
+
+class PILCaptureOutput(CaptureOutput):
+    def __init__(self) -> None: ...
+    def process(
+        self,
+        pointer: _CVoidConstPLike,
+        pitch: int,
+        size: int,
+        width: int,
+        height: int,
+        region: tuple[int, int, int, int],
+        rotation: int,
+    ) -> Image.Image: ...
+    def to_pil(self, frame: _ImageT) -> _ImageT: ...
+    def stack(self, frames: Sequence[_ImageT], stack_dimension: _Unused) -> Sequence[_ImageT]: ...

stubs/D3DShot/d3dshot/capture_outputs/pil_capture_output.pyi

@@ -0,0 +1,27 @@
+from _typeshed import Incomplete
+from collections.abc import Sequence
+from ctypes import _CVoidConstPLike
+from typing_extensions import Literal, TypeAlias
+
+from d3dshot.capture_output import CaptureOutput
+from PIL import Image
+
+# TODO: Complete types once we can import non-types dependencies
+# See: https://github.com/python/typeshed/issues/5768
+# from torch import Tensor
+_Tensor: TypeAlias = Incomplete
+
+class PytorchCaptureOutput(CaptureOutput):
+    def __init__(self) -> None: ...
+    def process(
+        self,
+        pointer: _CVoidConstPLike,
+        pitch: int,
+        size: int,
+        width: int,
+        height: int,
+        region: tuple[int, int, int, int],
+        rotation: int,
+    ) -> _Tensor: ...
+    def to_pil(self, frame: _Tensor) -> Image.Image: ...
+    def stack(self, frames: Sequence[_Tensor], stack_dimension: Literal["first", "last"]) -> _Tensor: ...

stubs/D3DShot/d3dshot/capture_outputs/pytorch_capture_output.pyi

@@ -0,0 +1,3 @@
+from d3dshot.capture_outputs.pytorch_capture_output import PytorchCaptureOutput
+
+class PytorchFloatCaptureOutput(PytorchCaptureOutput): ...

stubs/D3DShot/d3dshot/capture_outputs/pytorch_float_capture_output.pyi

@@ -0,0 +1,3 @@
+from d3dshot.capture_outputs.pytorch_gpu_capture_output import PytorchGPUCaptureOutput
+
+class PytorchFloatGPUCaptureOutput(PytorchGPUCaptureOutput): ...