typed-ffmpeg-compatible 2.1.0__py3-none-any.whl

typed_ffmpeg/dag/nodes.py

@@ -0,0 +1,479 @@
+from __future__ import annotations
+
+import logging
+import os.path
+from dataclasses import dataclass, replace
+from pathlib import Path
+from typing import TYPE_CHECKING, Any
+
+from ..exceptions import FFMpegTypeError, FFMpegValueError
+from ..schema import Default, StreamType
+from ..utils.escaping import escape
+from ..utils.lazy_eval.schema import LazyValue
+from ..utils.typing import override
+from .global_runnable.runnable import GlobalRunable
+from .io.output_args import OutputArgs
+from .schema import Node, Stream
+
+if TYPE_CHECKING:
+    from ..streams.audio import AudioStream
+    from ..streams.av import AVStream
+    from ..streams.video import VideoStream
+    from .context import DAGContext
+
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass(frozen=True, kw_only=True)
+class FilterNode(Node):
+    """
+    A filter node that can be used to apply filters to streams
+    """
+
+    name: str
+    """
+    The name of the filter
+    """
+
+    inputs: tuple[FilterableStream, ...] = ()
+    """
+    The input streams
+    """
+
+    input_typings: tuple[StreamType, ...] = ()
+    """
+    The input typings
+    """
+
+    output_typings: tuple[StreamType, ...] = ()
+    """
+    The output typings
+    """
+
+    @override
+    def repr(self) -> str:
+        return self.name
+
+    def video(self, index: int) -> "VideoStream":
+        """
+        Return the video stream at the specified index
+
+        Args:
+            index: the index of the video stream
+
+        Returns:
+            the video stream at the specified index
+        """
+        from ..streams.video import VideoStream
+
+        video_outputs = [i for i, k in enumerate(self.output_typings) if k == StreamType.video]
+        if not len(video_outputs) > index:
+            raise FFMpegValueError(f"Specified index {index} is out of range for video outputs {len(video_outputs)}")
+        return VideoStream(node=self, index=video_outputs[index])
+
+    def audio(self, index: int) -> "AudioStream":
+        """
+        Return the audio stream at the specified index
+
+        Args:
+            index: the index of the audio stream
+
+        Returns:
+            the audio stream at the specified index
+        """
+        from ..streams.audio import AudioStream
+
+        audio_outputs = [i for i, k in enumerate(self.output_typings) if k == StreamType.audio]
+        if not len(audio_outputs) > index:
+            raise FFMpegValueError(f"Specified index {index} is out of range for audio outputs {len(audio_outputs)}")
+
+        return AudioStream(node=self, index=audio_outputs[index])
+
+    def __post_init__(self) -> None:
+        from ..streams.audio import AudioStream
+        from ..streams.video import VideoStream
+
+        super().__post_init__()
+
+        if len(self.inputs) != len(self.input_typings):
+            raise FFMpegValueError(f"Expected {len(self.input_typings)} inputs, got {len(self.inputs)}")
+
+        stream: FilterableStream
+        expected_type: StreamType
+
+        for i, (stream, expected_type) in enumerate(zip(self.inputs, self.input_typings)):
+            if expected_type == StreamType.video:
+                if not isinstance(stream, VideoStream):
+                    raise FFMpegTypeError(
+                        f"Expected input {i} to have video component, got {stream.__class__.__name__}"
+                    )
+            if expected_type == StreamType.audio:
+                if not isinstance(stream, AudioStream):
+                    raise FFMpegTypeError(
+                        f"Expected input {i} to have audio component, got {stream.__class__.__name__}"
+                    )
+
+    @override
+    def get_args(self, context: DAGContext = None) -> list[str]:
+        from .context import DAGContext
+
+        if not context:
+            context = DAGContext.build(self)
+
+        incoming_labels = "".join(f"[{k.label(context)}]" for k in self.inputs)
+        outputs = context.get_outgoing_streams(self)
+
+        outgoing_labels = ""
+        for output in sorted(outputs, key=lambda stream: stream.index or 0):
+            # NOTE: all outgoing streams must be filterable
+            assert isinstance(output, FilterableStream)
+            outgoing_labels += f"[{output.label(context)}]"
+
+        commands = []
+        for key, value in self.kwargs:
+            assert not isinstance(value, LazyValue), f"LazyValue should have been evaluated: {key}={value}"
+
+            # Note: the -nooption syntax cannot be used for boolean AVOptions, use -option 0/-option 1.
+            if isinstance(value, bool):
+                value = str(int(value))
+
+            if not isinstance(value, Default):
+                commands += [f"{key}={escape(value)}"]
+
+        if commands:
+            return [incoming_labels] + [f"{self.name}="] + [escape(":".join(commands), "\\'[],;")] + [outgoing_labels]
+        return [incoming_labels] + [f"{self.name}"] + [outgoing_labels]
+
+
+@dataclass(frozen=True, kw_only=True)
+class FilterableStream(Stream, OutputArgs):
+    """
+    A stream that can be used as input to a filter
+    """
+
+    node: "FilterNode | InputNode"
+
+    @override
+    def _output_node(self, *streams: FilterableStream, filename: str | Path, **kwargs: Any) -> OutputNode:
+        """
+        Output the streams to a file URL
+
+        Args:
+            *streams: the other streams to output
+            filename: the filename to output to
+            **kwargs: the arguments for the output
+
+        Returns:
+            the output stream
+        """
+        return OutputNode(inputs=(self, *streams), filename=str(filename), kwargs=tuple(kwargs.items()))
+
+    def vfilter(
+        self,
+        *streams: "FilterableStream",
+        name: str,
+        input_typings: tuple[StreamType, ...] = (StreamType.video,),
+        **kwargs: Any,
+    ) -> "VideoStream":
+        """
+        Apply a custom video filter which has only one output to this stream
+
+        Args:
+            *streams (FilterableStream): the streams to apply the filter to
+            name: the name of the filter
+            input_typings: the input typings
+            **kwargs: the arguments for the filter
+
+        Returns:
+            the output stream
+        """
+        return self.filter_multi_output(
+            *streams,
+            name=name,
+            input_typings=input_typings,
+            output_typings=(StreamType.video,),
+            **kwargs,
+        ).video(0)
+
+    def afilter(
+        self,
+        *streams: "FilterableStream",
+        name: str,
+        input_typings: tuple[StreamType, ...] = (StreamType.audio,),
+        **kwargs: Any,
+    ) -> "AudioStream":
+        """
+        Apply a custom audio filter which has only one output to this stream
+
+        Args:
+            *streams (FilterableStream): the streams to apply the filter to
+            name: the name of the filter
+            input_typings: the input typings
+            **kwargs: the arguments for the filter
+
+        Returns:
+            the output stream
+        """
+        return self.filter_multi_output(
+            *streams,
+            name=name,
+            input_typings=input_typings,
+            output_typings=(StreamType.audio,),
+            **kwargs,
+        ).audio(0)
+
+    def filter_multi_output(
+        self,
+        *streams: "FilterableStream",
+        name: str,
+        input_typings: tuple[StreamType, ...] = (),
+        output_typings: tuple[StreamType, ...] = (),
+        **kwargs: Any,
+    ) -> "FilterNode":
+        """
+        Apply a custom filter which has multiple outputs to this stream
+
+        Args:
+            *streams (FilterableStream): the streams to apply the filter to
+            name: the name of the filter
+            input_typings: the input typings
+            output_typings: the output typings
+            **kwargs: the arguments for the filter
+
+        Returns:
+            the FilterNode
+        """
+        return FilterNode(
+            name=name,
+            kwargs=tuple(kwargs.items()),
+            inputs=(self, *streams),
+            input_typings=input_typings,
+            output_typings=output_typings,
+        )
+
+    def label(self, context: DAGContext = None) -> str:
+        """
+        Return the label for this stream
+
+        Args:
+            context: the DAG context
+
+        Returns:
+            the label for this stream
+        """
+        from ..streams.audio import AudioStream
+        from ..streams.av import AVStream
+        from ..streams.video import VideoStream
+        from .context import DAGContext
+
+        if not context:
+            context = DAGContext.build(self.node)
+
+        if isinstance(self.node, InputNode):
+            if isinstance(self, AVStream):
+                return f"{context.get_node_label(self.node)}"
+            elif isinstance(self, VideoStream):
+                return f"{context.get_node_label(self.node)}:v"
+            elif isinstance(self, AudioStream):
+                return f"{context.get_node_label(self.node)}:a"
+            raise FFMpegValueError(f"Unknown stream type: {self.__class__.__name__}")  # pragma: no cover
+
+        if isinstance(self.node, FilterNode):
+            if len(self.node.output_typings) > 1:
+                return f"{context.get_node_label(self.node)}#{self.index}"
+            return f"{context.get_node_label(self.node)}"
+        raise FFMpegValueError(f"Unknown node type: {self.node.__class__.__name__}")  # pragma: no cover
+
+    def __post_init__(self) -> None:
+        if isinstance(self.node, InputNode):
+            assert self.index is None, "Input streams cannot have an index"
+        else:
+            assert self.index is not None, "Filter streams must have an index"
+
+
+@dataclass(frozen=True, kw_only=True)
+class InputNode(Node):
+    """
+    A node that can be used to read from files
+    """
+
+    filename: str
+    """
+    The filename to read from
+    """
+
+    inputs: tuple[()] = ()
+
+    @override
+    def repr(self) -> str:
+        return os.path.basename(self.filename)
+
+    @property
+    def video(self) -> "VideoStream":
+        """
+        Return the video stream of this node
+
+        Returns:
+            the video stream
+        """
+        from ..streams.video import VideoStream
+
+        return VideoStream(node=self)
+
+    @property
+    def audio(self) -> "AudioStream":
+        """
+        Return the audio stream of this node
+
+        Returns:
+            the audio stream
+        """
+        from ..streams.audio import AudioStream
+
+        return AudioStream(node=self)
+
+    def stream(self) -> "AVStream":
+        """
+        Return the output stream of this node
+
+        Returns:
+            the output stream
+        """
+        from ..streams.av import AVStream
+
+        return AVStream(node=self)
+
+    @override
+    def get_args(self, context: DAGContext = None) -> list[str]:
+        commands = []
+        for key, value in self.kwargs:
+            if isinstance(value, bool):
+                if value is True:
+                    commands += [f"-{key}"]
+                # NOTE: the -nooption is not supported
+            else:
+                commands += [f"-{key}", str(value)]
+        commands += ["-i", self.filename]
+        return commands
+
+
+@dataclass(frozen=True, kw_only=True)
+class OutputNode(Node):
+    filename: str
+    """
+    The filename to output to
+    """
+    inputs: tuple[FilterableStream, ...]
+
+    @override
+    def repr(self) -> str:
+        return os.path.basename(self.filename)
+
+    def stream(self) -> "OutputStream":
+        """
+        Return the output stream of this node
+
+        Returns:
+            the output stream
+        """
+        return OutputStream(node=self)
+
+    @override
+    def get_args(self, context: DAGContext = None) -> list[str]:
+        # !handle mapping
+        commands = []
+
+        if context and (
+            any(isinstance(k.node, FilterNode) for k in self.inputs)
+            or len([k for k in context.all_nodes if isinstance(k, OutputNode)]) > 1
+        ):
+            for input in self.inputs:
+                if isinstance(input.node, InputNode):
+                    commands += ["-map", input.label(context)]
+                else:
+                    commands += ["-map", f"[{input.label(context)}]"]
+
+        for key, value in self.kwargs:
+            if isinstance(value, bool):
+                if value is True:
+                    commands += [f"-{key}"]
+                # NOTE: the -nooption is not supported
+            else:
+                commands += [f"-{key}", str(value)]
+        commands += [self.filename]
+        return commands
+
+
+@dataclass(frozen=True, kw_only=True)
+class OutputStream(Stream, GlobalRunable):
+    node: OutputNode
+
+    @override
+    def _global_node(self, *streams: OutputStream, **kwargs: Any) -> GlobalNode:
+        """
+        Add extra global command-line argument
+
+        Args:
+            **kwargs: the extra arguments
+
+        Returns:
+            the output stream
+        """
+        return GlobalNode(inputs=(self, *streams), kwargs=tuple(kwargs.items()))
+
+
+@dataclass(frozen=True, kw_only=True)
+class GlobalNode(Node):
+    """
+    A node that can be used to set global options
+    """
+
+    inputs: tuple[OutputStream, ...]
+
+    @override
+    def repr(self) -> str:
+        return " ".join(self.get_args())
+
+    def stream(self) -> "GlobalStream":
+        """
+        Return the output stream of this node
+
+        Returns:
+            the output stream
+        """
+        return GlobalStream(node=self)
+
+    @override
+    def get_args(self, context: DAGContext = None) -> list[str]:
+        commands = []
+        for key, value in self.kwargs:
+            if isinstance(value, bool):
+                if value is True:
+                    commands += [f"-{key}"]
+                # NOTE: the -no{key} format since not really for global options
+            else:
+                commands += [f"-{key}", str(value)]
+        return commands
+
+
+@dataclass(frozen=True, kw_only=True)
+class GlobalStream(Stream, GlobalRunable):
+    node: GlobalNode
+
+    @override
+    def _global_node(self, *streams: OutputStream, **kwargs: Any) -> GlobalNode:
+        """
+        Add extra global command-line argument
+
+        Args:
+            **kwargs: the extra arguments
+
+        Returns:
+            the output stream
+        """
+        inputs = (*self.node.inputs, *streams)
+        kwargs = dict(self.node.kwargs) | kwargs
+
+        new_node = replace(self.node, inputs=inputs, kwargs=tuple(kwargs.items()))
+        return new_node

typed_ffmpeg/dag/schema.py

@@ -0,0 +1,210 @@
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+from dataclasses import dataclass, replace
+from functools import cached_property
+from typing import TYPE_CHECKING, Literal
+
+from ..utils.lazy_eval.schema import LazyValue
+from .utils import is_dag
+
+if TYPE_CHECKING:
+    from .context import DAGContext
+
+
+@dataclass(frozen=True, kw_only=True)
+class HashableBaseModel:
+    """
+    A base class for hashable dataclasses.
+    """
+
+    @cached_property
+    def hex(self) -> str:
+        """
+        Get the hexadecimal hash of the object.
+        """
+        return hex(abs(hash(self)))[2:]
+
+
+@dataclass(frozen=True, kw_only=True)
+class Stream(HashableBaseModel):
+    """
+    A 'Stream' represents a sequence of data flow in the Directed Acyclic Graph (DAG).
+
+    Note:
+        Each stream in the DAG is a sequence of operations that transforms the data from its input form to its output form. The stream is an essential component of the DAG, as it defines the order and the nature of the operations that are performed on the data.
+    """
+
+    node: Node
+    """
+    Represents the node that the stream is connected to in the upstream direction.
+
+    Note:
+        In the context of a data stream, the 'upstream' refers to the source of the data, or where the data is coming from. Therefore, the 'upstream node' is the node that is providing the data to the current stream.
+    """
+
+    index: int | None = None  # the nth child of the node
+    """
+    Represents the index of the stream in the node's output streams.
+
+    Note:
+        See Also: [Stream specifiers](https://ffmpeg.org/ffmpeg.html#Stream-specifiers-1) `stream_index`
+    """
+
+    def view(self, format: Literal["png", "svg", "dot"] = "png") -> str:
+        """
+        Visualize the stream.
+
+        Args:
+            format: The format of the view.
+
+        Returns:
+            The file path of the visualization.
+        """
+        from ..utils.view import view
+
+        return view(self.node, format=format)
+
+    def _repr_png_(self) -> bytes:  # pragma: no cover
+        with open(self.view(format="png"), "rb") as f:
+            return f.read()
+
+    def _repr_svg_(self) -> str:  # pragma: no cover
+        with open(self.view(format="svg"), "r") as f:
+            return f.read()
+
+
+@dataclass(frozen=True, kw_only=True)
+class Node(HashableBaseModel, ABC):
+    """
+    A 'Node' represents a single operation in the Directed Acyclic Graph (DAG).
+
+    Note:
+        Each node in the DAG represents a single operation that transforms the data from its input form to its output form. The node is an essential component of the DAG, as it defines the nature of the operations that are performed on the data.
+    """
+
+    # Filter_Node_Option_Type
+    kwargs: tuple[tuple[str, str | int | float | bool | LazyValue], ...] = ()
+    """
+    Represents the keyword arguments of the node.
+    """
+
+    inputs: tuple[Stream, ...] = ()
+    """
+    Represents the input streams of the node.
+    """
+
+    def __post_init__(self) -> None:
+        # Validate the DAG
+        passed = set()
+        nodes = [self]
+        output = {}
+
+        while nodes:
+            node = nodes.pop()
+
+            if node in passed:
+                continue
+            passed.add(node)
+
+            nodes.extend(k.node for k in node.inputs)
+
+            output[node.hex] = set(k.node.hex for k in node.inputs)
+
+        if not is_dag(output):
+            raise ValueError(f"Graph is not a DAG: {output}")  # pragma: no cover
+
+    @abstractmethod
+    def get_args(self, context: DAGContext = None) -> list[str]:
+        """
+        Get the arguments of the node.
+
+        Args:
+            context: The DAG context.
+
+        Returns:
+            The arguments of the node.
+        """
+
+    def repr(self) -> str:
+        """
+        Get the representation of the node.
+
+        Returns:
+            The representation of the node.
+        """
+        return repr(self)
+
+    def replace(self, old_node: Node, new_node: Node) -> Node:
+        """
+        Replace the old node in the graph with the new node.
+
+        Args:
+            old_node: The old node to replace.
+            new_node: The new node to replace with.
+
+        Returns:
+            The new graph with the replaced node.
+        """
+        if self == old_node:
+            return new_node
+
+        inputs = []
+
+        for stream in self.inputs:
+            new_stream_node = stream.node.replace(old_node, new_node)
+
+            if new_stream_node != stream.node:
+                # need to create a new stream
+                new_stream = replace(stream, node=new_stream_node)
+                inputs.append(new_stream)
+            else:
+                inputs.append(stream)
+
+        return replace(self, inputs=tuple(inputs))
+
+    @property
+    def max_depth(self) -> int:
+        """
+        Get the maximum depth of the node.
+
+        Returns:
+            The maximum depth of the node.
+        """
+        return max((i.node.max_depth for i in self.inputs), default=0) + 1
+
+    @property
+    def upstream_nodes(self) -> set[Node]:
+        """
+        Get all upstream nodes of the node.
+
+        Returns:
+            The upstream nodes of the node.
+        """
+        output = {self}
+        for input in self.inputs:
+            output |= input.node.upstream_nodes
+
+        return output
+
+    def view(self, format: Literal["png", "svg", "dot"] = "png") -> str:
+        """
+        Visualize the Node.
+
+        Args:
+            format: The format of the view.
+
+        Returns:
+            The file path of the visualization.
+        """
+        from ..utils.view import view
+
+        return view(self, format=format)
+
+    def _repr_png_(self) -> bytes:  # pragma: no cover
+        with open(self.view(format="png"), "rb") as f:
+            return f.read()
+
+    def _repr_svg_(self) -> str:  # pragma: no cover
+        with open(self.view(format="svg"), "r") as f:
+            return f.read()