Added stubs for D3DShot

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/numpy_capture_output.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_float_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/pil_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/pytorch_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_float_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_gpu_capture_output.pyi

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

stubs/D3DShot/d3dshot/capture_outputs/pytorch_gpu_capture_output.pyi

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

stubs/D3DShot/d3dshot/d3dshot.pyi

@@ -0,0 +1,45 @@
+from collections import deque
+from collections.abc import Iterable
+
+from d3dshot.capture_output import CaptureOutput as CaptureOutput, CaptureOutputs as CaptureOutputs, _Frame
+from d3dshot.display import Display as Display
+
+class Singleton(type): ...
+
+class D3DShot(metaclass=Singleton):
+    displays: list[Display]
+    display: Display
+    capture_output: CaptureOutput
+    frame_buffer_size: int
+    frame_buffer: deque[_Frame]
+    previous_screenshot: _Frame | None
+    region: tuple[int, int, int, int] | None
+
+    def __init__(
+        self,
+        capture_output: CaptureOutputs = ...,
+        frame_buffer_size: int = ...,
+        pil_is_available: bool = ...,
+        numpy_is_available: bool = ...,
+        pytorch_is_available: bool = ...,
+        pytorch_gpu_is_available: bool = ...,
+    ) -> None: ...
+    @property
+    def is_capturing(self) -> bool: ...
+    def get_latest_frame(self) -> _Frame | None: ...
+    def get_frame(self, frame_index: int) -> _Frame | None: ...
+    def get_frames(self, frame_indices: Iterable[int]) -> list[_Frame]: ...
+    def get_frame_stack(self, frame_indices: Iterable[int], stack_dimension: str | None = ...) -> _Frame: ...
+    def screenshot(self, region: tuple[int, int, int, int] | None = ...) -> _Frame | None: ...
+    def screenshot_to_disk(
+        self, directory: str | None = ..., file_name: str | None = ..., region: tuple[int, int, int, int] | None = ...
+    ) -> str: ...
+    def frame_buffer_to_disk(self, directory: str | None = ...) -> None: ...
+    def capture(self, target_fps: int = ..., region: tuple[int, int, int, int] | None = ...) -> bool: ...
+    def screenshot_every(self, interval: float, region: tuple[int, int, int, int] | None = ...) -> bool: ...
+    def screenshot_to_disk_every(
+        self, interval: float, directory: str | None = ..., region: tuple[int, int, int, int] | None = ...
+    ) -> bool: ...
+    def stop(self) -> bool: ...
+    def benchmark(self) -> None: ...
+    def detect_displays(self) -> None: ...

stubs/D3DShot/d3dshot/display.pyi

@@ -0,0 +1,49 @@
+from ctypes import _Pointer
+from typing_extensions import TypedDict
+
+from d3dshot.dll import _ProcessFunc, _ProcessFuncRegionArg, _ProcessFuncReturn
+from d3dshot.dll.d3d import ID3D11Device, ID3D11DeviceContext
+from d3dshot.dll.dxgi import IDXGIAdapter, IDXGIOutput1, IDXGIOutputDuplication
+
+class _PositionDict(TypedDict):
+    left: int
+    top: int
+    right: int
+    bottom: int
+
+class Display:
+    name: str
+    adapter_name: str
+    resolution: tuple[int, int]
+    position: _PositionDict
+    rotation: int
+    scale_factor: float
+    is_primary: bool
+    hmonitor: int
+    dxgi_output: IDXGIOutput1 | None
+    dxgi_adapter: _Pointer[IDXGIAdapter] | None
+    # Note that Display.d3d_device and Display.d3d_device_context can never be None.
+    # Despite initially being set to None in __init__,
+    # they're always immediately set in _initialize_dxgi_output_duplication()
+    d3d_device: ID3D11Device
+    d3d_device_context: ID3D11DeviceContext
+    dxgi_output_duplication: _Pointer[IDXGIOutputDuplication]
+
+    def __init__(
+        self,
+        name: str | None = ...,
+        adapter_name: str | None = ...,
+        resolution: tuple[int, int] | None = ...,
+        position: _PositionDict | None = ...,
+        rotation: int | None = ...,
+        scale_factor: float | None = ...,
+        is_primary: bool = ...,
+        hmonitor: int | None = ...,
+        dxgi_output: IDXGIOutput1 | None = ...,
+        dxgi_adapter: _Pointer[IDXGIAdapter] | None = ...,
+    ) -> None: ...
+    def capture(
+        self, process_func: _ProcessFunc[_ProcessFuncRegionArg, _ProcessFuncReturn] | None, region: _ProcessFuncRegionArg = ...
+    ) -> _ProcessFuncReturn: ...
+    @classmethod
+    def discover_displays(cls) -> list[Display]: ...

stubs/D3DShot/d3dshot/dll/__init__.pyi

@@ -0,0 +1,28 @@
+import sys
+from _typeshed import Incomplete
+from collections.abc import Callable
+from ctypes import _CData, c_ulong
+from ctypes.wintypes import PFLOAT
+from typing import TypeVar
+from typing_extensions import TypeAlias
+
+from d3dshot.capture_output import _Frame
+
+_ProcessFuncRegionArg = TypeVar("_ProcessFuncRegionArg", tuple[int, int, int, int], None)
+_ProcessFuncReturn = TypeVar("_ProcessFuncReturn", _Frame, None)
+# The _ProcessFunc alias is used in multiple submodules
+_ProcessFunc: TypeAlias = Callable[[PFLOAT, int, int, int, int, _ProcessFuncRegionArg, int], _ProcessFuncReturn]  # noqa: Y047
+
+if sys.platform == "win32":
+    from ctypes import HRESULT
+
+    _HRESULT: TypeAlias = HRESULT
+else:
+    _HRESULT: TypeAlias = Incomplete
+
+# TODO: Use comtypes.IUnknown once we can import non-types dependencies
+# See: #5768
+class _IUnknown(_CData):
+    def QueryInterface(self, interface: type, iid: _CData | None = ...) -> _HRESULT: ...
+    def AddRef(self) -> c_ulong: ...
+    def Release(self) -> c_ulong: ...