Add correct stub

Author: mthrok

.github/workflows/_build_linux_cuda.yml

@@ -171,6 +171,8 @@ jobs:
 
       - name: Check Stub
         run: |
+          set -ex
+
           python scripts/generate_stub.py -o ~/stubs
           cp ~/stubs/* src/spdl/io/lib/
           git diff --exit-code
@@ -188,6 +190,8 @@ jobs:
       - name: Unit test
         if: '!cancelled()'
         run: |
+          set -ex
+
           # Run test
           python -c 'import spdl.io.utils;assert spdl.io.utils.built_with_cuda()'
           if "${{ inputs.use-nvdec == 'nvdec' }}" ; then

scripts/generate_stub.py

@@ -8,7 +8,6 @@
 
 """Generate interface files for extension modules."""
 
-import logging
 from argparse import ArgumentParser
 from pathlib import Path
 
@@ -31,7 +30,6 @@ def _generate(mod, output):
 
 def _main():
     args = _parse_args()
-    logging.basicConfig(level=logging.DEBUG)
     args.output_dir.mkdir(exist_ok=True, parents=True)
     _generate(_archive, args.output_dir / "_archive.pyi")
     _generate(_import_libspdl(), args.output_dir / "_libspdl.pyi")

src/spdl/io/lib/_libspdl_cuda.pyi

@@ -0,0 +1,58 @@
+from collections.abc import Callable, Sequence
+from typing import Annotated, overload
+
+from numpy.typing import NDArray
+
+import spdl.io.lib._spdl_ffmpeg8
+
+
+class CUDAConfig:
+    pass
+
+def cuda_config(device_index: int, stream: int = 0, allocator: tuple[Callable[[int, int, int], int], Callable[[int], None]] | None = None) -> CUDAConfig: ...
+
+class CUDABuffer:
+    @property
+    def __cuda_array_interface__(self) -> dict: ...
+
+    @property
+    def device_index(self) -> int: ...
+
+class NvDecDecoder:
+    def reset(self) -> None: ...
+
+    def init(self, device_config: CUDAConfig, codec: spdl.io.lib._spdl_ffmpeg8.VideoCodec, *, crop_left: int = 0, crop_top: int = 0, crop_right: int = 0, crop_bottom: int = 0, scale_width: int = -1, scale_height: int = -1) -> None: ...
+
+    def decode(self, packets: spdl.io.lib._spdl_ffmpeg8.VideoPackets) -> list[CUDABuffer]: ...
+
+    def flush(self) -> list[CUDABuffer]: ...
+
+@overload
+def decode_image_nvjpeg(data: bytes, *, device_config: CUDAConfig, scale_width: int = -1, scale_height: int = -1, pix_fmt: str = 'rgb', _zero_clear: bool = False) -> None: ...
+
+@overload
+def decode_image_nvjpeg(data: Sequence[bytes], *, device_config: CUDAConfig, scale_width: int, scale_height: int, pix_fmt: str = 'rgb', _zero_clear: bool = False) -> None: ...
+
+def cpu_storage(size: int) -> spdl.io.lib._spdl_ffmpeg8.CPUStorage: ...
+
+@overload
+def transfer_buffer(buffer: spdl.io.lib._spdl_ffmpeg8.CPUBuffer, *, device_config: CUDAConfig) -> CUDABuffer: ...
+
+@overload
+def transfer_buffer(buffer: Annotated[NDArray, dict(order='C', device='cpu')], *, device_config: CUDAConfig) -> CUDABuffer: ...
+
+def transfer_buffer_cpu(buffer: Annotated[NDArray, dict(order='C', device='cuda')]) -> spdl.io.lib._spdl_ffmpeg8.CPUBuffer: ...
+
+def init() -> None: ...
+
+def built_with_cuda() -> bool: ...
+
+def built_with_nvcodec() -> bool: ...
+
+def built_with_nvjpeg() -> bool: ...
+
+def synchronize_stream(arg: CUDAConfig, /) -> None: ...
+
+def nv12_to_planar_rgb(buffers: Sequence[CUDABuffer], *, device_config: CUDAConfig, matrix_coeff: int = 1) -> CUDABuffer: ...
+
+def nv12_to_planar_bgr(buffers: Sequence[CUDABuffer], *, device_config: CUDAConfig, matrix_coeff: int = 1) -> CUDABuffer: ...