vsjetengine 1.0.0__py3-none-any.whl

vsengine/policy.py

@@ -0,0 +1,395 @@
+# vs-engine
+# Copyright (C) 2022  cid-chan
+# Copyright (C) 2025  Jaded-Encoding-Thaumaturgy
+# This project is licensed under the EUPL-1.2
+# SPDX-License-Identifier: EUPL-1.2
+"""
+This module implements VapourSynth's Environment Policy system,
+allowing you to manage multiple VapourSynth environments within a single application.
+"""
+
+from __future__ import annotations
+
+import threading
+from abc import ABC, abstractmethod
+from collections.abc import Iterator
+from contextlib import AbstractContextManager, contextmanager
+from contextvars import ContextVar
+from logging import getLogger
+from types import MappingProxyType, TracebackType
+from typing import TYPE_CHECKING, Self
+from weakref import ReferenceType, ref
+
+from vapoursynth import (
+    AudioNode,
+    Core,
+    Environment,
+    EnvironmentData,
+    EnvironmentPolicy,
+    EnvironmentPolicyAPI,
+    VideoOutputTuple,
+    core,
+    get_outputs,
+    register_policy,
+)
+
+from ._hospice import admit_environment
+
+__all__ = ["ContextVarStore", "GlobalStore", "ManagedEnvironment", "Policy", "ThreadLocalStore"]
+
+
+logger = getLogger(__name__)
+
+
+class EnvironmentStore(ABC):
+    """
+    Environment Stores manage which environment is currently active.
+    """
+
+    @abstractmethod
+    def set_current_environment(self, environment: ReferenceType[EnvironmentData] | None) -> None:
+        """
+        Set the current environment in the store.
+        """
+
+    @abstractmethod
+    def get_current_environment(self) -> ReferenceType[EnvironmentData] | None:
+        """
+        Retrieve the current environment from the store (if any)
+        """
+
+
+class GlobalStore(EnvironmentStore):
+    """
+    This is the simplest store: It just stores the environment in a variable.
+    """
+
+    _current: ReferenceType[EnvironmentData] | None
+    __slots__ = ("_current",)
+
+    def set_current_environment(self, environment: ReferenceType[EnvironmentData] | None) -> None:
+        self._current = environment
+
+    def get_current_environment(self) -> ReferenceType[EnvironmentData] | None:
+        return getattr(self, "_current", None)
+
+
+class ThreadLocalStore(EnvironmentStore):
+    """
+    For simple threaded applications, use this store.
+
+    It will store the environment in a thread-local variable.
+    """
+
+    _current: threading.local
+
+    def __init__(self) -> None:
+        self._current = threading.local()
+
+    def set_current_environment(self, environment: ReferenceType[EnvironmentData] | None) -> None:
+        self._current.environment = environment
+
+    def get_current_environment(self) -> ReferenceType[EnvironmentData] | None:
+        return getattr(self._current, "environment", None)
+
+
+class ContextVarStore(EnvironmentStore):
+    """
+    If you are using AsyncIO or similar frameworks, use this store.
+    """
+
+    _current: ContextVar[ReferenceType[EnvironmentData] | None]
+
+    def __init__(self, name: str = "vapoursynth") -> None:
+        self._current = ContextVar(name)
+
+    def set_current_environment(self, environment: ReferenceType[EnvironmentData] | None) -> None:
+        self._current.set(environment)
+
+    def get_current_environment(self) -> ReferenceType[EnvironmentData] | None:
+        return self._current.get(None)
+
+
+class _ManagedPolicy(EnvironmentPolicy):
+    """
+    This class directly interfaces with VapourSynth.
+    """
+
+    __slots__ = ("_api", "_local", "_mutex", "_store")
+
+    def __init__(self, store: EnvironmentStore) -> None:
+        self._store = store
+        self._mutex = threading.Lock()
+        self._local = threading.local()
+
+    # For engine-calls that require vapoursynth but
+    # should not make their switch observable from the outside.
+
+    # Start the section.
+    def inline_section_start(self, environment: EnvironmentData) -> None:
+        self._local.environment = environment
+
+    # End the section.
+    def inline_section_end(self) -> None:
+        del self._local.environment
+
+    @property
+    def api(self) -> EnvironmentPolicyAPI:
+        if hasattr(self, "_api"):
+            return self._api
+
+        raise RuntimeError("Invalid state: No access to the current API")
+
+    def on_policy_registered(self, special_api: EnvironmentPolicyAPI) -> None:
+        self._api = special_api
+        logger.debug("Environment policy %r successfully registered with VapourSynth.", special_api)
+
+    def on_policy_cleared(self) -> None:
+        del self._api
+        logger.debug("Environment policy successfully cleared.")
+
+    def get_current_environment(self) -> EnvironmentData | None:
+        # For small segments, allow switching the environment inline.
+        # This is useful for vsengine-functions that require access to the
+        # vapoursynth api, but don't want to invoke the store for it.
+        if (env := getattr(self._local, "environment", None)) is not None and self.is_alive(env):
+            return env
+
+        # We wrap everything in a mutex to make sure
+        # no context-switch can reliably happen in this section.
+        with self._mutex:
+            current_environment = self._store.get_current_environment()
+            if current_environment is None:
+                return None
+
+            if current_environment() is None:
+                logger.warning("Environment reference from store resolved to dead object: %r", current_environment)
+                self._store.set_current_environment(None)
+                return None
+
+            received_environment = current_environment()
+
+            if TYPE_CHECKING:
+                assert received_environment
+
+            if not self.is_alive(received_environment):
+                logger.warning(
+                    "Received environment object is not alive (Garbage collected?): %r",
+                    received_environment,
+                )
+                # Remove the environment.
+                self._store.set_current_environment(None)
+                return None
+
+            return received_environment
+
+    def set_environment(self, environment: EnvironmentData | None) -> EnvironmentData | None:
+        with self._mutex:
+            previous_environment = self._store.get_current_environment()
+
+            if environment is not None and not self.is_alive(environment):
+                logger.warning("Attempted to set environment which is not alive: %r", environment)
+                self._store.set_current_environment(None)
+            else:
+                logger.debug("Environment successfully set to: %r", environment)
+                if environment is None:
+                    self._store.set_current_environment(None)
+                else:
+                    self._store.set_current_environment(ref(environment))
+
+            if previous_environment is not None:
+                return previous_environment()
+
+        return None
+
+
+class ManagedEnvironment(AbstractContextManager["ManagedEnvironment"]):
+    """
+    Represents a VapourSynth environment that is managed by a policy.
+    """
+
+    __slots__ = ("_data", "_environment", "_policy")
+
+    def __init__(self, environment: Environment, data: EnvironmentData, policy: Policy) -> None:
+        self._environment = environment
+        self._data = data
+        self._policy = policy
+
+    def __enter__(self) -> Self:
+        return self
+
+    def __exit__(self, exc: type[BaseException] | None, val: BaseException | None, tb: TracebackType | None) -> None:
+        self.dispose()
+
+    @property
+    def core(self) -> Core:
+        """
+        Returns the core representing this environment.
+        """
+        with self.inline_section():
+            return core.core
+
+    @property
+    def disposed(self) -> bool:
+        """
+        Checks if the environment is disposed
+        """
+        return not hasattr(self, "_data")
+
+    @property
+    def outputs(self) -> MappingProxyType[int, VideoOutputTuple | AudioNode]:
+        """
+        Returns the outputs within this environment.
+        """
+        with self.inline_section():
+            return get_outputs()
+
+    @property
+    def vs_environment(self) -> Environment:
+        """
+        Returns the vapoursynth.Environment-object representing this environment.
+        """
+        return self._environment
+
+    def switch(self) -> None:
+        """
+        Switches to the given environment without storing
+        which environment has been defined previously.
+        """
+        self._environment.use().__enter__()
+
+    def dispose(self) -> None:
+        """
+        Disposes of the environment.
+        """
+        if self.disposed:
+            return
+
+        logger.debug("Starting disposal of environment: %r", self._data)
+
+        admit_environment(self._data, self.core)
+        self._policy.api.destroy_environment(self._data)
+        del self._data
+
+    @contextmanager
+    def inline_section(self) -> Iterator[None]:
+        """
+        Private API!
+
+        Switches to the given environment within the block without
+        notifying the store.
+
+        If you follow the rules below, switching the environment
+        will be invisible to the caller.
+
+        Rules for safely calling this function:
+        - Do not suspend greenlets within the block!
+        - Do not yield or await within the block!
+        - Do not use __enter__ and __exit__ directly.
+        - This function is not reentrant.
+        """
+        self._policy.managed.inline_section_start(self._data)
+        try:
+            yield
+        finally:
+            self._policy.managed.inline_section_end()
+
+    @contextmanager
+    def use(self) -> Iterator[None]:
+        """
+        Switches to this environment within a block.
+        """
+        # prev_environment = self._policy.managed._store.get_current_environment()
+        with self._environment.use():
+            yield
+
+        # FIXME
+        # # Workaround: On 32bit systems, environment policies do not reset.
+        # self._policy.managed.set_environment(prev_environment)
+
+    def __del__(self) -> None:
+        if self.disposed:
+            return
+
+        import warnings
+
+        warnings.warn(f"Disposing {self!r} inside __del__. This might cause leaks.", ResourceWarning)
+        self.dispose()
+
+
+class Policy(AbstractContextManager["Policy"]):
+    """
+    A managed policy is a very simple policy that just stores the environment
+    data within the given store.
+
+    For convenience (especially for testing), this is a context manager that
+    makes sure policies are being unregistered when leaving a block.
+    """
+
+    _managed: _ManagedPolicy
+
+    def __init__(self, store: EnvironmentStore, flags_creation: int = 0) -> None:
+        """
+        Initializes a new Policy
+
+        Args:
+            store: The store to use for managing environments.
+            flags_creation: The flags to use when creating environments.
+                            See vapoursynth.CoreCreationFlags for more information.
+        """
+        self._managed = _ManagedPolicy(store)
+        self.flags_creation = flags_creation
+
+    def __enter__(self) -> Self:
+        self.register()
+        return self
+
+    def __exit__(self, _: type[BaseException] | None, __: BaseException | None, ___: TracebackType | None) -> None:
+        self.unregister()
+
+    @property
+    def api(self) -> EnvironmentPolicyAPI:
+        """
+        Returns the API instance for more complex interactions.
+
+        You will rarely need to use this directly.
+        """
+        return self._managed.api
+
+    @property
+    def managed(self) -> _ManagedPolicy:
+        """
+        Returns the actual policy within VapourSynth.
+
+        You will rarely need to use this directly.
+        """
+        return self._managed
+
+    def register(self) -> None:
+        """
+        Registers the policy with VapourSynth.
+        """
+        register_policy(self._managed)
+
+    def unregister(self) -> None:
+        """
+        Unregisters the policy from VapourSynth.
+        """
+        self._managed.api.unregister_policy()
+
+    def new_environment(self) -> ManagedEnvironment:
+        """
+        Creates a new VapourSynth core.
+
+        You need to call `dispose()` on this environment when you are done
+        using the new environment.
+
+        For convenience, a managed environment will also serve as a
+        context-manager that disposes the environment automatically.
+        """
+        data = self.api.create_environment(self.flags_creation)
+        env = self.api.wrap_environment(data)
+
+        menv = ManagedEnvironment(env, data, self)
+        logger.debug("Successfully created new environment %r", data)
+        return menv

vsengine/video.py

@@ -0,0 +1,180 @@
+# vs-engine
+# Copyright (C) 2022  cid-chan
+# Copyright (C) 2025  Jaded-Encoding-Thaumaturgy
+# This project is licensed under the EUPL-1.2
+# SPDX-License-Identifier: EUPL-1.2
+"""
+vsengine.render renders video frames for you.
+"""
+
+from collections.abc import Iterator, Sequence
+from concurrent.futures import Future
+
+import vapoursynth as vs
+
+from ._futures import UnifiedFuture, unified
+from ._helpers import use_inline
+from ._nodes import buffer_futures, close_when_needed
+from .policy import ManagedEnvironment
+
+__all__ = ["frame", "frames", "planes", "render"]
+
+
+@unified(kind="future")
+def frame(
+    node: vs.VideoNode, frameno: int, env: vs.Environment | ManagedEnvironment | None = None
+) -> Future[vs.VideoFrame]:
+    """
+    Request a specific frame from a node.
+
+    :param node: The node to request the frame from.
+    :param frameno: The frame number to request.
+    :param env: The environment to use for the request.
+    :return: A future that resolves to the frame.
+    """
+    with use_inline("frame", env):
+        return node.get_frame_async(frameno)
+
+
+@unified(kind="future")
+def planes(
+    node: vs.VideoNode,
+    frameno: int,
+    env: vs.Environment | ManagedEnvironment | None = None,
+    *,
+    planes: Sequence[int] | None = None,
+) -> Future[tuple[bytes, ...]]:
+    """
+    Request a specific frame from a node and return the planes as bytes.
+
+    :param node: The node to request the frame from.
+    :param frameno: The frame number to request.
+    :param env: The environment to use for the request.
+    :param planes: The planes to return. If None, all planes are returned.
+    :return: A future that resolves to a tuple of bytes.
+    """
+
+    def _extract(frame: vs.VideoFrame) -> tuple[bytes, ...]:
+        try:
+            # This might be a variable format clip.
+            # extract the plane as late as possible.
+            ps = range(len(frame)) if planes is None else planes
+            return tuple(bytes(frame[p]) for p in ps)
+        finally:
+            frame.close()
+
+    return frame(node, frameno, env).map(_extract)
+
+
+@unified(kind="generator")
+def frames(
+    node: vs.VideoNode,
+    env: vs.Environment | ManagedEnvironment | None = None,
+    *,
+    prefetch: int = 0,
+    backlog: int | None = None,
+    # Unlike the implementation provided by VapourSynth,
+    # we don't have to care about backwards compatibility and
+    # can just do the right thing from the beginning.
+    close: bool = True,
+) -> Iterator[Future[vs.VideoFrame]]:
+    """
+    Iterate over the frames of a node.
+
+    :param node: The node to iterate over.
+    :param env: The environment to use for the request.
+    :param prefetch: The number of frames to prefetch.
+    :param backlog: The maximum number of frames to keep in the backlog.
+    :param close: Whether to close the frames automatically.
+    :return: An iterator of futures that resolve to the frames.
+    """
+    with use_inline("frames", env):
+        length = len(node)
+
+    it = (frame(node, n, env) for n in range(length))
+
+    # If backlog is zero, skip.
+    if backlog is None or backlog > 0:
+        it = buffer_futures(it, prefetch=prefetch, backlog=backlog)
+
+    if close:
+        it = close_when_needed(it)
+
+    return it
+
+
+@unified(kind="generator")
+def render(
+    node: vs.VideoNode,
+    env: vs.Environment | ManagedEnvironment | None = None,
+    *,
+    prefetch: int = 0,
+    backlog: int | None = 0,
+    y4m: bool = False,
+) -> Iterator[Future[tuple[int, bytes]]]:
+    """
+    Render a node to a stream of bytes.
+
+    :param node: The node to render.
+    :param env: The environment to use for the request.
+    :param prefetch: The number of frames to prefetch.
+    :param backlog: The maximum number of frames to keep in the backlog.
+    :param y4m: Whether to output a Y4M header.
+    :return: An iterator of futures that resolve to a tuple of the frame number and the frame data.
+    """
+    frame_count = len(node)
+
+    if y4m:
+        match node.format.color_family:
+            case vs.GRAY:
+                y4mformat = "mono"
+            case vs.YUV:
+                match (node.format.subsampling_w, node.format.subsampling_h):
+                    case (1, 1):
+                        y4mformat = "420"
+                    case (1, 0):
+                        y4mformat = "422"
+                    case (0, 0):
+                        y4mformat = "444"
+                    case (2, 2):
+                        y4mformat = "410"
+                    case (2, 0):
+                        y4mformat = "411"
+                    case (0, 1):
+                        y4mformat = "440"
+                    case _:
+                        raise NotImplementedError
+            case _:
+                raise ValueError("Can only use GRAY and YUV for Y4M-Streams")
+
+        if node.format.bits_per_sample > 8:
+            y4mformat += f"p{node.format.bits_per_sample}"
+
+        y4mformat = "C" + y4mformat + " "
+
+        data = "YUV4MPEG2 {y4mformat}W{width} H{height} F{fps_num}:{fps_den} Ip A0:0 XLENGTH={length}\n".format(  # noqa: UP032
+            y4mformat=y4mformat,
+            width=node.width,
+            height=node.height,
+            fps_num=node.fps_num,
+            fps_den=node.fps_den,
+            length=frame_count,
+        )
+        yield UnifiedFuture.resolve((0, data.encode("ascii")))
+
+    current_frame = 0
+
+    def render_single_frame(frame: vs.VideoFrame) -> tuple[int, bytes]:
+        buf = list[bytes]()
+
+        if y4m:
+            buf.append(b"FRAME\n")
+
+        for plane in iter(frame):
+            buf.append(bytes(plane))
+
+        return current_frame, b"".join(buf)
+
+    for frame, fut in enumerate(frames(node, env, prefetch=prefetch, backlog=backlog).futures, 1):
+        current_frame = frame
+        yield UnifiedFuture.from_future(fut).map(render_single_frame)