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: ...