capturekit 0.1.0__py3-none-any.whl

capturekit/_pipeline/frame_queue.py

@@ -0,0 +1,120 @@
+"""Bounded frame queue with duplicate-slot merging, plus a buffer pool.
+
+Only *new* frames hold buffers and count against the capacity. A duplicate
+slot ("repeat the previous frame") is a counter merged into the queue tail,
+so the pacer never blocks and memory stays bounded even when the encoder
+falls behind.
+"""
+
+from __future__ import annotations
+
+import collections
+import threading
+from dataclasses import dataclass
+
+import numpy as np
+import numpy.typing as npt
+
+Buffer = npt.NDArray[np.uint8]
+
+
+@dataclass(slots=True)
+class FrameItem:
+    buffer: Buffer
+    timestamp_ns: int
+
+
+@dataclass(slots=True)
+class DupItem:
+    count: int
+
+
+QueueItem = FrameItem | DupItem
+
+
+class FrameQueue:
+    def __init__(self, capacity: int) -> None:
+        self._capacity = capacity
+        self._items: collections.deque[QueueItem] = collections.deque()
+        self._frames = 0
+        self._closed = False
+        self._cond = threading.Condition()
+        self.max_depth = 0
+
+    def put_frame(self, item: FrameItem) -> bool:
+        """Enqueue a new frame. Returns False (not enqueued) if full."""
+        with self._cond:
+            if self._closed:
+                return False
+            if self._frames >= self._capacity:
+                return False
+            self._items.append(item)
+            self._frames += 1
+            self.max_depth = max(self.max_depth, self._frames)
+            self._cond.notify()
+            return True
+
+    def put_dup(self, count: int = 1) -> None:
+        with self._cond:
+            if self._closed or count <= 0:
+                return
+            if self._items and isinstance(self._items[-1], DupItem):
+                self._items[-1].count += count
+            else:
+                self._items.append(DupItem(count))
+            self._cond.notify()
+
+    def get(self, timeout_s: float | None = None) -> QueueItem | None:
+        """Next item, or ``None`` when closed and drained (or timed out)."""
+        with self._cond:
+            if not self._cond.wait_for(lambda: self._items or self._closed, timeout=timeout_s):
+                return None
+            if not self._items:
+                return None
+            item = self._items.popleft()
+            if isinstance(item, FrameItem):
+                self._frames -= 1
+            return item
+
+    def close(self) -> None:
+        """No more puts; consumers drain remaining items then get ``None``."""
+        with self._cond:
+            self._closed = True
+            self._cond.notify_all()
+
+    def drain(self) -> list[QueueItem]:
+        with self._cond:
+            items = list(self._items)
+            self._items.clear()
+            self._frames = 0
+            return items
+
+    @property
+    def depth(self) -> int:
+        with self._cond:
+            return self._frames
+
+
+class BufferPool:
+    """Recycles fixed-shape frame buffers to avoid per-frame large allocations."""
+
+    def __init__(self, shape: tuple[int, int, int], max_buffers: int) -> None:
+        self.shape = shape
+        self._max = max_buffers
+        self._free: list[Buffer] = []
+        self._allocated = 0
+        self._lock = threading.Lock()
+
+    def acquire(self) -> Buffer | None:
+        """A free buffer, or ``None`` if ``max_buffers`` are all in use."""
+        with self._lock:
+            if self._free:
+                return self._free.pop()
+            if self._allocated >= self._max:
+                return None
+            self._allocated += 1
+        return np.zeros(self.shape, dtype=np.uint8)
+
+    def release(self, buf: Buffer) -> None:
+        with self._lock:
+            self._free.append(buf)

capturekit/_pipeline/mixer.py

@@ -0,0 +1,80 @@
+"""Mixes several aligned sample streams into one."""
+
+from __future__ import annotations
+
+import threading
+from collections import deque
+from collections.abc import Callable
+
+import numpy as np
+import numpy.typing as npt
+
+Samples = npt.NDArray[np.float32]
+
+
+class SampleFifo:
+    def __init__(self, channels: int) -> None:
+        self.channels = channels
+        self._blocks: deque[Samples] = deque()
+        self.size = 0
+
+    def push(self, block: Samples) -> None:
+        if block.shape[0]:
+            self._blocks.append(block)
+            self.size += int(block.shape[0])
+
+    def pop(self, n: int) -> Samples:
+        out = np.empty((n, self.channels), dtype=np.float32)
+        filled = 0
+        while filled < n:
+            blk = self._blocks[0]
+            take = min(n - filled, blk.shape[0])
+            out[filled : filled + take] = blk[:take]
+            filled += take
+            if take == blk.shape[0]:
+                self._blocks.popleft()
+            else:
+                self._blocks[0] = blk[take:]
+        self.size -= n
+        return out
+
+
+class AudioMixer:
+    """Sums N aligned inputs; emits only samples available on all live inputs.
+
+    Inputs marked finished count as silence so one ended source cannot stall
+    the mix. Output is clipped to [-1, 1].
+    """
+
+    def __init__(self, inputs: int, channels: int, sink: Callable[[Samples], None]) -> None:
+        self._fifos = [SampleFifo(channels) for _ in range(inputs)]
+        self._finished = [False] * inputs
+        self._sink = sink
+        self._lock = threading.Lock()
+        self.channels = channels
+
+    def push(self, index: int, block: Samples) -> None:
+        with self._lock:
+            self._fifos[index].push(block)
+            self._emit_locked(final=False)
+
+    def finish_input(self, index: int) -> None:
+        with self._lock:
+            self._finished[index] = True
+            self._emit_locked(final=all(self._finished))
+
+    def _emit_locked(self, final: bool) -> None:
+        live = [f.size for f, done in zip(self._fifos, self._finished, strict=True) if not done]
+        if live and not final:
+            n = min(live)
+        else:
+            n = max((f.size for f in self._fifos), default=0)
+        if n <= 0:
+            return
+        mix = np.zeros((n, self.channels), dtype=np.float32)
+        for f in self._fifos:
+            take = min(n, f.size)
+            if take:
+                mix[:take] += f.pop(take)
+        np.clip(mix, -1.0, 1.0, out=mix)
+        self._sink(mix)

capturekit/_pipeline/timeline.py

@@ -0,0 +1,133 @@
+"""Pause-aware mapping between wall-clock time and media time."""
+
+from __future__ import annotations
+
+import threading
+from dataclasses import dataclass
+
+
+@dataclass(slots=True)
+class _Segment:
+    wall_start: int
+    media_start: int
+    wall_end: int | None = None
+
+
+class Timeline:
+    """Maps capturekit-clock nanoseconds to media nanoseconds.
+
+    Media time starts at 0 when :meth:`start` is called and does not advance
+    while paused. All methods are thread-safe.
+    """
+
+    def __init__(self) -> None:
+        self._lock = threading.Lock()
+        self._cond = threading.Condition(self._lock)
+        self._segments: list[_Segment] = []
+        self._stopped = False
+
+    # -- state changes -------------------------------------------------
+    def start(self, wall_ns: int) -> None:
+        with self._cond:
+            if self._segments:
+                raise RuntimeError("timeline already started")
+            self._segments.append(_Segment(wall_ns, 0))
+            self._cond.notify_all()
+
+    def pause(self, wall_ns: int) -> int:
+        """Pause at ``wall_ns``; returns the media time at the pause."""
+        with self._cond:
+            seg = self._active()
+            if seg is None:
+                raise RuntimeError("timeline not running")
+            wall_ns = max(wall_ns, seg.wall_start)
+            seg.wall_end = wall_ns
+            return seg.media_start + (wall_ns - seg.wall_start)
+
+    def resume(self, wall_ns: int) -> None:
+        with self._cond:
+            if not self._segments or self._segments[-1].wall_end is None:
+                raise RuntimeError("timeline not paused")
+            last = self._segments[-1]
+            assert last.wall_end is not None
+            media = last.media_start + (last.wall_end - last.wall_start)
+            self._segments.append(_Segment(max(wall_ns, last.wall_end), media))
+            self._cond.notify_all()
+
+    def stop(self, wall_ns: int) -> int:
+        """Stop at ``wall_ns`` (pausing if running); returns final media time."""
+        with self._cond:
+            self._stopped = True
+            self._cond.notify_all()
+            if not self._segments:
+                return 0
+            seg = self._segments[-1]
+            if seg.wall_end is None:
+                seg.wall_end = max(wall_ns, seg.wall_start)
+            return seg.media_start + (seg.wall_end - seg.wall_start)
+
+    # -- queries -------------------------------------------------------
+    def _active(self) -> _Segment | None:
+        if self._segments and self._segments[-1].wall_end is None:
+            return self._segments[-1]
+        return None
+
+    @property
+    def generation(self) -> int:
+        """Number of recording segments so far (increments on every resume)."""
+        with self._lock:
+            return len(self._segments)
+
+    @property
+    def started(self) -> bool:
+        with self._lock:
+            return bool(self._segments)
+
+    @property
+    def paused(self) -> bool:
+        return self.is_paused()
+
+    def is_paused(self) -> bool:
+        """Method form of :attr:`paused` (re-evaluated on each call; another
+        thread may pause or resume at any time)."""
+        with self._lock:
+            return bool(self._segments) and self._segments[-1].wall_end is not None
+
+    @property
+    def stopped(self) -> bool:
+        with self._lock:
+            return self._stopped
+
+    def media_ns(self, wall_ns: int) -> int | None:
+        """Media time of ``wall_ns`` or ``None`` if it falls outside recording."""
+        with self._lock:
+            for seg in reversed(self._segments):
+                if wall_ns >= seg.wall_start:
+                    if seg.wall_end is not None and wall_ns >= seg.wall_end:
+                        return None
+                    return seg.media_start + (wall_ns - seg.wall_start)
+            return None
+
+    def media_now(self, wall_now: int) -> int:
+        """Media time reached at ``wall_now`` (clamped while paused)."""
+        with self._lock:
+            if not self._segments:
+                return 0
+            seg = self._segments[-1]
+            end = wall_now if seg.wall_end is None else min(wall_now, seg.wall_end)
+            return seg.media_start + max(0, end - seg.wall_start)
+
+    def wall_for_media(self, media_ns: int) -> int | None:
+        """Wall time at which the *active* segment reaches ``media_ns``."""
+        with self._lock:
+            seg = self._active()
+            if seg is None:
+                return None
+            return seg.wall_start + (media_ns - seg.media_start)
+
+    def wait_running(self, timeout_s: float) -> bool:
+        """Block until the timeline is running (not paused) or stopped."""
+        with self._cond:
+            return self._cond.wait_for(
+                lambda: self._stopped or self._active() is not None, timeout=timeout_s
+            )

capturekit/_platform.py

@@ -0,0 +1,24 @@
+"""Lazy selection of the platform backend module."""
+
+from __future__ import annotations
+
+import sys
+from types import ModuleType
+
+from .errors import UnsupportedPlatformError
+
+
+def backend() -> ModuleType:
+    """Return the platform ``api`` module (imported on first use)."""
+    plat: str = sys.platform  # a variable, so type checkers keep every branch reachable
+    if plat == "win32":
+        from ._backends.windows import api as win_api
+
+        return win_api
+    if plat.startswith("linux"):
+        from ._backends.linux import api as linux_api
+
+        return linux_api
+    raise UnsupportedPlatformError(
+        f"capturekit supports Windows 10/11 and Linux X11; {sys.platform!r} is not supported"
+    )

capturekit/_version.py

@@ -0,0 +1 @@
+__version__ = "0.1.0"

capturekit/_worker.py

@@ -0,0 +1,100 @@
+"""A single thread that owns a video backend and executes calls on it."""
+
+from __future__ import annotations
+
+import queue
+import threading
+from collections.abc import Callable
+from concurrent.futures import Future
+from concurrent.futures import TimeoutError as FuturesTimeout
+from typing import Any, TypeVar
+
+from . import _platform
+from ._backends.base import VideoBackend
+from .errors import CaptureError, CaptureTimeoutError
+from .sources import VideoSource
+
+T = TypeVar("T")
+
+
+class BackendThread:
+    """Creates, starts and drives a backend on one dedicated thread.
+
+    Native video APIs (COM/WinRT on Windows, an Xlib ``Display`` on Linux) are
+    used from this thread only; callers interact via :meth:`call`.
+    """
+
+    def __init__(self, source: VideoSource, name: str = "capturekit-video") -> None:
+        self._source = source
+        self._calls: queue.Queue[tuple[Callable[[VideoBackend], Any], Future[Any]] | None] = (
+            queue.Queue()
+        )
+        self._ready: Future[VideoBackend] = Future()
+        self._thread = threading.Thread(target=self._run, name=name, daemon=True)
+        self._closed = threading.Event()
+
+    def start(self, timeout_s: float = 15.0) -> VideoBackend:
+        self._thread.start()
+        try:
+            return self._ready.result(timeout_s)
+        except FuturesTimeout:
+            # The backend may still finish starting later: make it exit and
+            # close itself instead of leaking a live capture session.
+            self._calls.put(None)
+            raise CaptureTimeoutError(
+                f"capture backend did not start within {timeout_s}s"
+            ) from None
+
+    def _run(self) -> None:
+        api = _platform.backend()
+        try:
+            with api.thread_context():
+                backend: VideoBackend = api.open_video(self._source)
+                try:
+                    backend.start()
+                except BaseException as exc:
+                    backend.close()
+                    self._ready.set_exception(exc)
+                    return
+                self._ready.set_result(backend)
+                try:
+                    while True:
+                        item = self._calls.get()
+                        if item is None:
+                            break
+                        fn, fut = item
+                        if not fut.set_running_or_notify_cancel():
+                            continue
+                        try:
+                            fut.set_result(fn(backend))
+                        except BaseException as exc:
+                            fut.set_exception(exc)
+                finally:
+                    backend.close()
+        except BaseException as exc:
+            if not self._ready.done():
+                self._ready.set_exception(exc)
+        finally:
+            self._closed.set()
+            # Never leave a caller waiting on a call that will not run.
+            while True:
+                try:
+                    pending = self._calls.get_nowait()
+                except queue.Empty:
+                    break
+                if pending is not None and not pending[1].done():
+                    pending[1].set_exception(CaptureError("capture was closed"))
+
+    def call(self, fn: Callable[[VideoBackend], T], timeout_s: float | None = None) -> T:
+        if self._closed.is_set() or not self._thread.is_alive():
+            raise CaptureError("capture is closed")
+        fut: Future[T] = Future()
+        self._calls.put((fn, fut))
+        if self._closed.is_set() and not fut.done():  # raced with shutdown
+            fut.set_exception(CaptureError("capture was closed"))
+        return fut.result(timeout_s)
+
+    def stop(self, timeout_s: float = 10.0) -> None:
+        if self._thread.is_alive():
+            self._calls.put(None)
+            self._thread.join(timeout_s)

capturekit/capabilities.py

@@ -0,0 +1,79 @@
+"""Programmatic capability discovery."""
+
+from __future__ import annotations
+
+import enum
+from dataclasses import dataclass, field
+
+
+class Support(enum.Enum):
+    #: Implemented and tested on this platform.
+    SUPPORTED = "supported"
+    #: Implemented; limited testing or known gaps (see ``reason``).
+    EXPERIMENTAL = "experimental"
+    #: Implemented for this platform, but this machine/OS build/server lacks
+    #: a requirement (``reason`` says which).
+    UNAVAILABLE = "unavailable"
+    #: Not implemented or out of scope.
+    UNSUPPORTED = "unsupported"
+
+
+@dataclass(frozen=True, slots=True)
+class Capability:
+    status: Support
+    reason: str = ""
+
+    def __bool__(self) -> bool:
+        return self.status in (Support.SUPPORTED, Support.EXPERIMENTAL)
+
+
+def supported(reason: str = "") -> Capability:
+    return Capability(Support.SUPPORTED, reason)
+
+
+def experimental(reason: str) -> Capability:
+    return Capability(Support.EXPERIMENTAL, reason)
+
+
+def unavailable(reason: str) -> Capability:
+    return Capability(Support.UNAVAILABLE, reason)
+
+
+def unsupported(reason: str) -> Capability:
+    return Capability(Support.UNSUPPORTED, reason)
+
+
+@dataclass(frozen=True, slots=True)
+class Capabilities:
+    platform: str
+    os_version: str
+    display_server: str | None
+    monitor_capture: Capability
+    window_capture: Capability
+    window_capture_borderless: Capability
+    window_capture_minimized: Capability
+    cursor_capture: Capability
+    microphone_audio: Capability
+    system_audio: Capability
+    process_audio: Capability
+    audio_session_enumeration: Capability
+    recording: Capability
+    encoders: tuple[str, ...] = ()
+    notes: tuple[str, ...] = field(default_factory=tuple)
+
+    def summary(self) -> str:
+        lines = [
+            f"capturekit on {self.platform} {self.os_version} "
+            f"(display server: {self.display_server or 'n/a'})"
+        ]
+        for name in (
+            "monitor_capture", "window_capture", "window_capture_borderless",
+            "window_capture_minimized", "cursor_capture", "microphone_audio", "system_audio",
+            "process_audio", "audio_session_enumeration", "recording",
+        ):  # fmt: skip
+            cap: Capability = getattr(self, name)
+            lines.append(f"  {name:28s} {cap.status.value:12s} {cap.reason}")
+        if self.encoders:
+            lines.append(f"  validated encoders: {', '.join(self.encoders)}")
+        lines.extend(f"  note: {n}" for n in self.notes)
+        return "\n".join(lines)

capturekit/capture.py

@@ -0,0 +1,141 @@
+"""Frame capture without recording."""
+
+from __future__ import annotations
+
+import threading
+from collections.abc import Iterator
+from types import TracebackType
+
+import numpy as np
+
+from ._backends.base import AcquiredFrame, VideoBackend
+from ._pipeline.clock import NS_PER_S, now_ns, sleep_until
+from ._worker import BackendThread
+from .errors import CaptureTimeoutError
+from .sources import MonitorSource, VideoSource, WindowSource
+from .types import Frame, Monitor, Rect, Window
+
+
+class ScreenCapture:
+    """Capture frames from a monitor or a window.
+
+    Example::
+
+        with ScreenCapture(MonitorSource(0)) as cap:
+            frame = cap.grab()          # newest frame (BGRA numpy array)
+            for frame in cap.frames(fps=30, max_frames=90):
+                ...
+
+    Thread-safe; all native work happens on an internal backend thread.
+    """
+
+    def __init__(self, source: VideoSource) -> None:
+        self.source = source
+        self._worker: BackendThread | None = None
+        self._backend_name = ""
+        self._last: Frame | None = None
+        self._index = 0
+        self._lock = threading.Lock()
+
+    # -- lifecycle -----------------------------------------------------------
+    def open(self) -> ScreenCapture:
+        with self._lock:
+            self._open_locked()
+        return self
+
+    def _open_locked(self) -> None:
+        if self._worker is None:
+            worker = BackendThread(self.source, name="capturekit-capture")
+            backend = worker.start()
+            self._backend_name = backend.info.name
+            self._worker = worker
+
+    def close(self) -> None:
+        with self._lock:
+            if self._worker is not None:
+                self._worker.stop()
+                self._worker = None
+
+    def __enter__(self) -> ScreenCapture:
+        return self.open()
+
+    def __exit__(
+        self, et: type[BaseException] | None, ev: BaseException | None, tb: TracebackType | None
+    ) -> None:
+        self.close()
+
+    @property
+    def backend_name(self) -> str:
+        return self._backend_name
+
+    # -- capture ---------------------------------------------------------------
+    def _acquire(self, timeout_s: float) -> AcquiredFrame | None:
+        if self._worker is None:
+            self._open_locked()
+        assert self._worker is not None
+
+        def fn(b: VideoBackend) -> AcquiredFrame | None:
+            return b.acquire(None, timeout_s)
+
+        return self._worker.call(fn, timeout_s + 10)
+
+    def grab(self, timeout: float = 2.0, *, new_only: bool = False) -> Frame:
+        """Return the newest frame.
+
+        If nothing changed since the previous call, a copy of the previous
+        frame is returned (unless ``new_only``). Raises
+        :class:`CaptureTimeoutError` if no frame is available in ``timeout``.
+        """
+        with self._lock:
+            first_wait = timeout if (self._last is None or new_only) else 0.0
+            acq = self._acquire(first_wait)
+            if acq is None:
+                if self._last is None or new_only:
+                    raise CaptureTimeoutError(f"no frame within {timeout:.2f}s")
+                last = self._last
+                return Frame(last.data.copy(), last.timestamp_ns, last.index, last.content_size)
+            frame = Frame(acq.data, acq.timestamp_ns, self._index, acq.content_size)
+            self._index += 1
+            self._last = frame
+            return Frame(frame.data.copy(), frame.timestamp_ns, frame.index, frame.content_size)
+
+    def frames(
+        self, fps: float | None = 30.0, *, max_frames: int | None = None, timeout: float = 2.0
+    ) -> Iterator[Frame]:
+        """Iterate frames.
+
+        With ``fps`` frames are produced at absolute deadlines (repeating the
+        previous image when nothing changed). With ``fps=None`` only new
+        frames are yielded, as fast as the source produces them.
+        """
+        count = 0
+        if fps is None:
+            while max_frames is None or count < max_frames:
+                yield self.grab(timeout, new_only=True)
+                count += 1
+            return
+        period = NS_PER_S / fps
+        t0 = now_ns()
+        n = 0
+        while max_frames is None or count < max_frames:
+            sleep_until(t0 + int(n * period))
+            yield self.grab(timeout)
+            count += 1
+            # Skip deadlines already missed (slow consumer) instead of bursting.
+            n = max(n + 1, int((now_ns() - t0) / period))
+
+
+def screenshot(
+    target: Monitor | Window | int | None = None,
+    region: Rect | None = None,
+    *,
+    timeout: float = 2.0,
+) -> np.ndarray:
+    """One-shot BGRA screenshot of a monitor (default primary) or window."""
+    src: VideoSource
+    if isinstance(target, Window):
+        src = WindowSource(target, region=region)
+    else:
+        src = MonitorSource(target, region=region)
+    with ScreenCapture(src) as cap:
+        return cap.grab(timeout).data