Skip to content

Commit 3ec7b1f

Browse files
AvasamAvasam
authored
Add stubs for D3DShot (#8652)
1 parent 34383fc commit 3ec7b1f

18 files changed

+664
-0
lines changed

‎stubs/D3DShot/METADATA.toml

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

‎stubs/D3DShot/d3dshot/__init__.pyi

+12
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,12 @@
1+
from d3dshot.capture_output import CaptureOutputs as CaptureOutputs
2+
from d3dshot.d3dshot import D3DShot as D3DShot
3+
4+
pil_is_available: bool
5+
numpy_is_available: bool
6+
pytorch_is_available: bool
7+
pytorch_gpu_is_available: bool
8+
capture_output_mapping: dict[str, CaptureOutputs]
9+
capture_outputs: list[str]
10+
11+
def determine_available_capture_outputs() -> list[CaptureOutputs]: ...
12+
def create(capture_output: str = ..., frame_buffer_size: int = ...) -> D3DShot: ...

‎stubs/D3DShot/d3dshot/capture_output.pyi

+44
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,44 @@
1+
import enum
2+
from _typeshed import Incomplete
3+
from collections.abc import Sequence
4+
from ctypes import _CVoidConstPLike
5+
from typing_extensions import Literal, TypeAlias
6+
7+
from PIL import Image
8+
9+
_Frame: TypeAlias = Image.Image | Incomplete
10+
# TODO: Complete types once we can import non-types dependencies
11+
# See: #5768
12+
# from torch import Tensor
13+
# from comtypes import IUnknown
14+
# import numpy.typing as npt
15+
# _Frame: TypeAlias = Image.Image | npt.NDArray[np.int32] | npt.NDArray[np.float32] | Tensor
16+
17+
class CaptureOutputs(enum.Enum):
18+
PIL: int
19+
NUMPY: int
20+
NUMPY_FLOAT: int
21+
PYTORCH: int
22+
PYTORCH_FLOAT: int
23+
PYTORCH_GPU: int
24+
PYTORCH_FLOAT_GPU: int
25+
26+
class CaptureOutputError(BaseException): ...
27+
28+
# All CaptureOutput methods just reference the backend. Making this both a base class and a wrapper.
29+
class CaptureOutput:
30+
# `backend` is a subclass of CaptureOutput based on the CaptureOutputs enum passed to __init__
31+
backend: CaptureOutput
32+
def __init__(self, backend: CaptureOutputs = ...) -> None: ...
33+
def process(
34+
self,
35+
pointer: _CVoidConstPLike,
36+
pitch: int,
37+
size: int,
38+
width: int,
39+
height: int,
40+
region: tuple[int, int, int, int],
41+
rotation: int,
42+
) -> _Frame: ...
43+
def to_pil(self, frame: _Frame) -> Image.Image: ...
44+
def stack(self, frames: Sequence[_Frame], stack_dimension: Literal["first", "last"]) -> _Frame: ...

‎stubs/D3DShot/d3dshot/capture_outputs/__init__.pyi

Whitespace-only changes.

‎stubs/D3DShot/d3dshot/capture_outputs/numpy_capture_output.pyi

+29
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,29 @@
1+
from _typeshed import Incomplete
2+
from collections.abc import Sequence
3+
from ctypes import _CVoidConstPLike
4+
from typing_extensions import Literal, TypeAlias
5+
6+
from d3dshot.capture_output import CaptureOutput
7+
from PIL import Image
8+
9+
# TODO: Complete types once we can import non-types dependencies
10+
# See: #5768
11+
# import numpy as np
12+
# import numpy.typing as npt
13+
# _NDArray: TypeAlias = npt.NDArray[np.int32]
14+
_NDArray: TypeAlias = Incomplete
15+
16+
class NumpyCaptureOutput(CaptureOutput):
17+
def __init__(self) -> None: ...
18+
def process(
19+
self,
20+
pointer: _CVoidConstPLike,
21+
pitch: int,
22+
size: int,
23+
width: int,
24+
height: int,
25+
region: tuple[int, int, int, int],
26+
rotation: int,
27+
) -> _NDArray: ...
28+
def to_pil(self, frame: _NDArray) -> Image.Image: ...
29+
def stack(self, frames: Sequence[_NDArray] | _NDArray, stack_dimension: Literal["first", "last"]) -> _NDArray: ...

‎stubs/D3DShot/d3dshot/capture_outputs/numpy_float_capture_output.pyi

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

‎stubs/D3DShot/d3dshot/capture_outputs/pil_capture_output.pyi

+25
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,25 @@
1+
from collections.abc import Sequence
2+
from ctypes import _CVoidConstPLike
3+
from typing import TypeVar
4+
from typing_extensions import TypeAlias
5+
6+
from d3dshot.capture_output import CaptureOutput
7+
from PIL import Image
8+
9+
_Unused: TypeAlias = object
10+
_ImageT = TypeVar("_ImageT", bound=Image.Image)
11+
12+
class PILCaptureOutput(CaptureOutput):
13+
def __init__(self) -> None: ...
14+
def process(
15+
self,
16+
pointer: _CVoidConstPLike,
17+
pitch: int,
18+
size: int,
19+
width: int,
20+
height: int,
21+
region: tuple[int, int, int, int],
22+
rotation: int,
23+
) -> Image.Image: ...
24+
def to_pil(self, frame: _ImageT) -> _ImageT: ...
25+
def stack(self, frames: Sequence[_ImageT], stack_dimension: _Unused) -> Sequence[_ImageT]: ...

‎stubs/D3DShot/d3dshot/capture_outputs/pytorch_capture_output.pyi

+27
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,27 @@
1+
from _typeshed import Incomplete
2+
from collections.abc import Sequence
3+
from ctypes import _CVoidConstPLike
4+
from typing_extensions import Literal, TypeAlias
5+
6+
from d3dshot.capture_output import CaptureOutput
7+
from PIL import Image
8+
9+
# TODO: Complete types once we can import non-types dependencies
10+
# See: https://github.com/python/typeshed/issues/5768
11+
# from torch import Tensor
12+
_Tensor: TypeAlias = Incomplete
13+
14+
class PytorchCaptureOutput(CaptureOutput):
15+
def __init__(self) -> None: ...
16+
def process(
17+
self,
18+
pointer: _CVoidConstPLike,
19+
pitch: int,
20+
size: int,
21+
width: int,
22+
height: int,
23+
region: tuple[int, int, int, int],
24+
rotation: int,
25+
) -> _Tensor: ...
26+
def to_pil(self, frame: _Tensor) -> Image.Image: ...
27+
def stack(self, frames: Sequence[_Tensor], stack_dimension: Literal["first", "last"]) -> _Tensor: ...

‎stubs/D3DShot/d3dshot/capture_outputs/pytorch_float_capture_output.pyi

+3
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,3 @@
1+
from d3dshot.capture_outputs.pytorch_capture_output import PytorchCaptureOutput
2+
3+
class PytorchFloatCaptureOutput(PytorchCaptureOutput): ...

‎stubs/D3DShot/d3dshot/capture_outputs/pytorch_float_gpu_capture_output.pyi

+3
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,3 @@
1+
from d3dshot.capture_outputs.pytorch_gpu_capture_output import PytorchGPUCaptureOutput
2+
3+
class PytorchFloatGPUCaptureOutput(PytorchGPUCaptureOutput): ...

‎stubs/D3DShot/d3dshot/capture_outputs/pytorch_gpu_capture_output.pyi

+3
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,3 @@
1+
from d3dshot.capture_outputs.pytorch_capture_output import PytorchCaptureOutput
2+
3+
class PytorchGPUCaptureOutput(PytorchCaptureOutput): ...

‎stubs/D3DShot/d3dshot/d3dshot.pyi

+45
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,45 @@
1+
from collections import deque
2+
from collections.abc import Iterable
3+
4+
from d3dshot.capture_output import CaptureOutput as CaptureOutput, CaptureOutputs as CaptureOutputs, _Frame
5+
from d3dshot.display import Display as Display
6+
7+
class Singleton(type): ...
8+
9+
class D3DShot(metaclass=Singleton):
10+
displays: list[Display]
11+
display: Display
12+
capture_output: CaptureOutput
13+
frame_buffer_size: int
14+
frame_buffer: deque[_Frame]
15+
previous_screenshot: _Frame | None
16+
region: tuple[int, int, int, int] | None
17+
18+
def __init__(
19+
self,
20+
capture_output: CaptureOutputs = ...,
21+
frame_buffer_size: int = ...,
22+
pil_is_available: bool = ...,
23+
numpy_is_available: bool = ...,
24+
pytorch_is_available: bool = ...,
25+
pytorch_gpu_is_available: bool = ...,
26+
) -> None: ...
27+
@property
28+
def is_capturing(self) -> bool: ...
29+
def get_latest_frame(self) -> _Frame | None: ...
30+
def get_frame(self, frame_index: int) -> _Frame | None: ...
31+
def get_frames(self, frame_indices: Iterable[int]) -> list[_Frame]: ...
32+
def get_frame_stack(self, frame_indices: Iterable[int], stack_dimension: str | None = ...) -> _Frame: ...
33+
def screenshot(self, region: tuple[int, int, int, int] | None = ...) -> _Frame | None: ...
34+
def screenshot_to_disk(
35+
self, directory: str | None = ..., file_name: str | None = ..., region: tuple[int, int, int, int] | None = ...
36+
) -> str: ...
37+
def frame_buffer_to_disk(self, directory: str | None = ...) -> None: ...
38+
def capture(self, target_fps: int = ..., region: tuple[int, int, int, int] | None = ...) -> bool: ...
39+
def screenshot_every(self, interval: float, region: tuple[int, int, int, int] | None = ...) -> bool: ...
40+
def screenshot_to_disk_every(
41+
self, interval: float, directory: str | None = ..., region: tuple[int, int, int, int] | None = ...
42+
) -> bool: ...
43+
def stop(self) -> bool: ...
44+
def benchmark(self) -> None: ...
45+
def detect_displays(self) -> None: ...

‎stubs/D3DShot/d3dshot/display.pyi

+49
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,49 @@
1+
from ctypes import _Pointer
2+
from typing_extensions import TypedDict
3+
4+
from d3dshot.dll import _ProcessFunc, _ProcessFuncRegionArg, _ProcessFuncReturn
5+
from d3dshot.dll.d3d import ID3D11Device, ID3D11DeviceContext
6+
from d3dshot.dll.dxgi import IDXGIAdapter, IDXGIOutput1, IDXGIOutputDuplication
7+
8+
class _PositionDict(TypedDict):
9+
left: int
10+
top: int
11+
right: int
12+
bottom: int
13+
14+
class Display:
15+
name: str
16+
adapter_name: str
17+
resolution: tuple[int, int]
18+
position: _PositionDict
19+
rotation: int
20+
scale_factor: float
21+
is_primary: bool
22+
hmonitor: int
23+
dxgi_output: IDXGIOutput1 | None
24+
dxgi_adapter: _Pointer[IDXGIAdapter] | None
25+
# Note that Display.d3d_device and Display.d3d_device_context can never be None.
26+
# Despite initially being set to None in __init__,
27+
# they're always immediately set in _initialize_dxgi_output_duplication()
28+
d3d_device: ID3D11Device
29+
d3d_device_context: ID3D11DeviceContext
30+
dxgi_output_duplication: _Pointer[IDXGIOutputDuplication]
31+
32+
def __init__(
33+
self,
34+
name: str | None = ...,
35+
adapter_name: str | None = ...,
36+
resolution: tuple[int, int] | None = ...,
37+
position: _PositionDict | None = ...,
38+
rotation: int | None = ...,
39+
scale_factor: float | None = ...,
40+
is_primary: bool = ...,
41+
hmonitor: int | None = ...,
42+
dxgi_output: IDXGIOutput1 | None = ...,
43+
dxgi_adapter: _Pointer[IDXGIAdapter] | None = ...,
44+
) -> None: ...
45+
def capture(
46+
self, process_func: _ProcessFunc[_ProcessFuncRegionArg, _ProcessFuncReturn] | None, region: _ProcessFuncRegionArg = ...
47+
) -> _ProcessFuncReturn: ...
48+
@classmethod
49+
def discover_displays(cls) -> list[Display]: ...

‎stubs/D3DShot/d3dshot/dll/__init__.pyi

+28
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,28 @@
1+
import sys
2+
from _typeshed import Incomplete
3+
from collections.abc import Callable
4+
from ctypes import _CData, c_ulong
5+
from ctypes.wintypes import PFLOAT
6+
from typing import TypeVar
7+
from typing_extensions import TypeAlias
8+
9+
from d3dshot.capture_output import _Frame
10+
11+
_ProcessFuncRegionArg = TypeVar("_ProcessFuncRegionArg", tuple[int, int, int, int], None)
12+
_ProcessFuncReturn = TypeVar("_ProcessFuncReturn", _Frame, None)
13+
# The _ProcessFunc alias is used in multiple submodules
14+
_ProcessFunc: TypeAlias = Callable[[PFLOAT, int, int, int, int, _ProcessFuncRegionArg, int], _ProcessFuncReturn] # noqa: Y047
15+
16+
if sys.platform == "win32":
17+
from ctypes import HRESULT
18+
19+
_HRESULT: TypeAlias = HRESULT
20+
else:
21+
_HRESULT: TypeAlias = Incomplete
22+
23+
# TODO: Use comtypes.IUnknown once we can import non-types dependencies
24+
# See: #5768
25+
class _IUnknown(_CData):
26+
def QueryInterface(self, interface: type, iid: _CData | None = ...) -> _HRESULT: ...
27+
def AddRef(self) -> c_ulong: ...
28+
def Release(self) -> c_ulong: ...

0 commit comments

Comments
 (0)