sera-2 1.17.0__py3-none-any.whl → 1.19.0__py3-none-any.whl

sera/libs/directed_computing_graph/__init__.py

@@ -0,0 +1,53 @@
+"""
+Directed Computing Graph package for sera.
+
+This package provides classes and utilities for working with directed computing graphs.
+"""
+
+# Import type aliases and annotated types
+# Import all classes from submodules
+from ._dcg import DirectedComputingGraph, NodeId, TaskArgs, TaskKey
+from ._edge import DCGEdge
+from ._flow import Flow
+from ._node import ComputeFn, ComputeFnId, DCGNode, PartialFn
+from ._runtime import SKIP, UNSET, ArgValueType, NodeRuntime
+
+# Import utility functions from type conversion
+from ._type_conversion import (
+    ComposeTypeConversion,
+    TypeConversion,
+    UnitTypeConversion,
+    align_generic_type,
+    ground_generic_type,
+    is_generic_type,
+    patch_get_origin,
+)
+
+# Define __all__ to control what gets exported
+__all__ = [
+    # Main classes
+    "DirectedComputingGraph",
+    "DCGNode",
+    "DCGEdge",
+    "Flow",
+    "PartialFn",
+    "TypeConversion",
+    "NodeRuntime",
+    # Enums and special values
+    "ArgValueType",
+    "UNSET",
+    "SKIP",
+    # Type aliases and annotations
+    "NodeId",
+    "TaskKey",
+    "TaskArgs",
+    "ComputeFnId",
+    "ComputeFn",
+    "UnitTypeConversion",
+    "ComposeTypeConversion",
+    # Utility functions
+    "patch_get_origin",
+    "is_generic_type",
+    "align_generic_type",
+    "ground_generic_type",
+]

sera/libs/directed_computing_graph/_dcg.py

@@ -0,0 +1,403 @@
+from __future__ import annotations
+
+import asyncio
+import inspect
+from dataclasses import dataclass
+from enum import Enum
+from typing import (
+    Annotated,
+    Any,
+    Awaitable,
+    Callable,
+    MutableSequence,
+    Optional,
+    Sequence,
+)
+
+from graph.retworkx import RetworkXStrDiGraph
+
+from sera.libs.directed_computing_graph._edge import DCGEdge
+from sera.libs.directed_computing_graph._flow import Flow
+from sera.libs.directed_computing_graph._node import (
+    ComputeFn,
+    ComputeFnId,
+    DCGNode,
+    NodeId,
+)
+from sera.libs.directed_computing_graph._runtime import SKIP, NodeRuntime
+from sera.libs.directed_computing_graph._type_conversion import (
+    ComposeTypeConversion,
+    TypeConversion,
+    UnitTypeConversion,
+    align_generic_type,
+    ground_generic_type,
+    is_generic_type,
+)
+from sera.misc import identity
+
+TaskKey = Annotated[tuple, "TaskKey"]
+TaskArgs = Annotated[MutableSequence, "TaskArgs"]
+
+
+class DirectedComputingGraph:
+    """
+    A Directed Computing Graph (DCG) is a directed graph where nodes represent functions
+    and edges represent dependencies between these functions. The graph is used to manage
+    the execution of functions in a specific order based on their dependencies.
+    """
+
+    def __init__(
+        self,
+        graph: RetworkXStrDiGraph[int, DCGNode, DCGEdge],
+        type_service: TypeConversion,
+    ):
+        self.graph = graph
+        self.type_service = type_service
+
+    @staticmethod
+    def from_flows(
+        flows: dict[ComputeFnId, Flow | ComputeFn],
+        type_conversions: Optional[
+            Sequence[UnitTypeConversion | ComposeTypeConversion]
+        ] = None,
+        strict: bool = True,
+    ):
+        """Create a computing graph from flow mapping.
+
+        Args:
+            flows: A dictionary mapping identifier to:
+                1. a function
+                2. a flow specifying the upstream functions and the function.
+            type_conversions: A list of type conversions to be used for converting the input types.
+            strict: If True, we do type checking.
+        Returns:
+            DirectedComputingGraph: A directed computing graph constructed from the provided flows.
+        """
+        # add typing conversions
+        upd_type_conversions: list[UnitTypeConversion | ComposeTypeConversion] = list(
+            type_conversions or []
+        )
+        type_service = TypeConversion(upd_type_conversions)
+
+        g: RetworkXStrDiGraph[int, DCGNode, DCGEdge] = RetworkXStrDiGraph(
+            check_cycle=False, multigraph=False
+        )
+
+        # create a graph
+        for uid, uinfo in flows.items():
+            if isinstance(uinfo, Flow):
+                actor = uinfo.target
+            else:
+                actor = uinfo
+            g.add_node(DCGNode(uid, actor))
+
+        # create a graph
+        for uid, uinfo in flows.items():
+            if isinstance(uinfo, Flow):
+                func = uinfo.target
+            else:
+                func = uinfo
+            g.add_node(DCGNode(uid, func))
+
+        # grounding function that has generic type input and output
+        for uid, flow in flows.items():
+            if not isinstance(flow, Flow):
+                continue
+
+            u = g.get_node(uid)
+            usig = u.signature
+            if is_generic_type(usig.return_type) or any(
+                is_generic_type(t) for t in usig.argtypes
+            ):
+                var2type = {}
+                for i, t in enumerate(usig.argtypes):
+                    if is_generic_type(t):
+                        # align the generic type with the previous return type
+                        if len(flow.source) <= i and strict:
+                            raise TypeConversion.UnknownConversion(
+                                f"Cannot ground the generic type based on upstream actors for actor {uid}"
+                            )
+
+                        source_return_type = g.get_node(
+                            flow.source[i]
+                        ).signature.return_type
+
+                        try:
+                            usig.argtypes[i], (var, nt) = align_generic_type(
+                                t, source_return_type
+                            )
+                        except Exception as e:
+                            raise TypeConversion.UnknownConversion(
+                                f"Cannot align the generic type {t} based on upstream actors for actor {uid}"
+                            )
+                        var2type[var] = nt
+                if is_generic_type(usig.return_type):
+                    usig.return_type = ground_generic_type(
+                        usig.return_type,
+                        var2type,
+                    )
+
+        for uid, flow in flows.items():
+            if not isinstance(flow, Flow):
+                continue
+
+            u = g.get_node(uid)
+            usig = u.signature
+            for idx, sid in enumerate(flow.source):
+                s = g.get_node(sid)
+                ssig = s.signature
+                cast_fn = identity
+                try:
+                    cast_fn = type_service.get_conversion(
+                        ssig.return_type, usig.argtypes[idx]
+                    )
+                except Exception as e:
+                    if strict:
+                        raise TypeConversion.UnknownConversion(
+                            f"Don't know how to convert output of `{sid}` to input of `{uid}`"
+                        ) from e
+                g.add_edge(
+                    DCGEdge(
+                        id=-1,
+                        source=sid,
+                        target=uid,
+                        argindex=idx,
+                        filter_fn=flow.filter_fn,
+                        type_conversion=cast_fn,
+                    )
+                )
+
+        # postprocessing such as type conversion, and args/context
+        for u in g.iter_nodes():
+            inedges = g.in_edges(u.id)
+
+            # update the type conversion
+            u.type_conversions = [identity] * len(u.signature.argnames)
+            for inedge in inedges:
+                u.type_conversions[inedge.argindex] = inedge.type_conversion
+
+            # update the required args and context
+            u.required_args = u.signature.argnames[: g.in_degree(u.id)]
+            # arguments of a compute function that are not provided by the upstream actors must be provided by the context.
+            u.required_context = u.signature.argnames[g.in_degree(u.id) :]
+            u.required_context_default_args = {
+                k: u.signature.default_args[k]
+                for k in u.required_context
+                if k in u.signature.default_args
+            }
+
+        return DirectedComputingGraph(g, type_service)
+
+    def execute(
+        self,
+        input: dict[ComputeFnId, tuple],
+        output: set[str],
+        context: Optional[
+            dict[str, Callable | Any] | Callable[[], dict[str, Any]]
+        ] = None,
+    ):
+        """
+        Execute the directed computing graph with the given input and output specifications.
+
+        Args:
+            input: A dictionary mapping function identifiers to their input arguments.
+            output: A set of function identifiers that should be executed.
+            context: An optional context that can be a dictionary of functions or a single function.
+        """
+        assert all(
+            isinstance(v, tuple) for v in input.values()
+        ), "Input must be a tuple"
+
+        if context is None:
+            context = {}
+        elif isinstance(context, Callable):
+            context = context()
+        else:
+            context = {k: v() if callable(v) else v for k, v in context.items()}
+
+        # This is a quick reactive algorithm, we may be able to do it better.
+        # The idea is when all inputs of a function is available, we can execute a function.
+        # We assume that the memory is large enough to hold all the functions and their inputs
+        # in the memory.
+
+        # we execute the computing nodes
+        # when it's finished, we put the outgoing edges into a stack.
+        runtimes: dict[NodeId, NodeRuntime] = {}
+
+        for u in self.graph.iter_nodes():
+            if u.id in input:
+                # user provided input should supersede the context
+                n_provided_args = len(input[u.id])
+                n_consumed_context = n_provided_args - len(u.required_args)
+            else:
+                n_consumed_context = 0
+
+            node_context = tuple(
+                (
+                    context[name]
+                    if name in context
+                    else u.required_context_default_args[name]
+                )
+                for name in u.required_context[n_consumed_context:]
+            )
+
+            runtimes[u.id] = NodeRuntime.from_node(self.graph, u, node_context)
+        stack: list[NodeId] = []
+
+        for id, args in input.items():
+            runtimes[id].add_task((0,), list(args))
+            stack.append(id)
+
+        return_output = {id: [] for id in output}
+
+        while len(stack) > 0:
+            # pop the one from the stack and execute it.
+            id = stack.pop()
+            runtime = runtimes[id]
+
+            # if there is enough data for the node, we can execute it.
+            # if it is not, we just skip it and it will be added back to the stack by one of its parents.
+            # so we don't miss it.
+            if not runtime.has_enough_data():
+                continue
+
+            outedges = self.graph.out_edges(id)
+            successors: Sequence[tuple[DCGEdge, DCGNode]] = [
+                (edge, self.graph.get_node(edge.target)) for edge in outedges
+            ]
+
+            # run the tasks and pass the output to the successors
+            for task_id, task in runtime.tasks.items():
+                if any(arg is SKIP for arg in task):
+                    task_output = SKIP
+                else:
+                    task_output = runtime.execute(task)
+
+                for outedge, succ in successors:
+                    runtimes[succ.id].add_task_args(
+                        task_id,
+                        id,
+                        (
+                            SKIP
+                            if task_output is SKIP or not outedge.filter(task_output)
+                            else task_output
+                        ),
+                    )
+
+                if id in output and task_output is not SKIP:
+                    return_output[id].append(task_output)
+
+            # retrieve the outgoing nodes and push them into the stack
+            for outedge, succ in successors:
+                stack.append(succ.id)
+
+        return return_output
+
+    async def execute_async(
+        self,
+        input: dict[ComputeFnId, tuple],
+        output: set[str],
+        context: Optional[
+            dict[str, Callable | Any] | Callable[[], dict[str, Any]]
+        ] = None,
+    ):
+        """
+        Asynchronously execute the directed computing graph with the given input and output specifications.
+        This method handles both synchronous and asynchronous functions.
+
+        Args:
+            input: A dictionary mapping function identifiers to their input arguments.
+            output: A set of function identifiers that should be executed.
+            context: An optional context that can be a dictionary of functions or a single function.
+        """
+        assert all(
+            isinstance(v, tuple) for v in input.values()
+        ), "Input must be a tuple"
+
+        if context is None:
+            context = {}
+        elif isinstance(context, Callable):
+            context = context()
+        else:
+            context = {k: v() if callable(v) else v for k, v in context.items()}
+
+        # This is a quick reactive algorithm, we may be able to do it better.
+        # The idea is when all inputs of a function is available, we can execute a function.
+        # We assume that the memory is large enough to hold all the functions and their inputs
+        # in the memory.
+
+        # we execute the computing nodes
+        # when it's finished, we put the outgoing edges into a stack.
+        runtimes: dict[NodeId, NodeRuntime] = {}
+
+        for u in self.graph.iter_nodes():
+            if u.id in input:
+                # user provided input should supersede the context
+                n_provided_args = len(input[u.id])
+                n_consumed_context = n_provided_args - len(u.required_args)
+            else:
+                n_consumed_context = 0
+
+            node_context = tuple(
+                (
+                    context[name]
+                    if name in context
+                    else u.required_context_default_args[name]
+                )
+                for name in u.required_context[n_consumed_context:]
+            )
+
+            runtimes[u.id] = NodeRuntime.from_node(self.graph, u, node_context)
+        stack: list[NodeId] = []
+
+        for id, args in input.items():
+            runtimes[id].add_task((0,), list(args))
+            stack.append(id)
+
+        return_output = {id: [] for id in output}
+
+        while len(stack) > 0:
+            # pop the one from the stack and execute it.
+            id = stack.pop()
+            runtime = runtimes[id]
+
+            # if there is enough data for the node, we can execute it.
+            # if it is not, we just skip it and it will be added back to the stack by one of its parents.
+            # so we don't miss it.
+            if not runtime.has_enough_data():
+                continue
+
+            outedges = self.graph.out_edges(id)
+            successors: Sequence[tuple[DCGEdge, DCGNode]] = [
+                (edge, self.graph.get_node(edge.target)) for edge in outedges
+            ]
+
+            # run the tasks and pass the output to the successors
+            for task_id, task in runtime.tasks.items():
+                if any(arg is SKIP for arg in task):
+                    task_output = SKIP
+                else:
+                    if runtime.node.signature.is_async:
+                        task_output = await runtime.execute(task)
+                    else:
+                        task_output = runtime.execute(task)
+
+                for outedge, succ in successors:
+                    runtimes[succ.id].add_task_args(
+                        task_id,
+                        id,
+                        (
+                            SKIP
+                            if task_output is SKIP or not outedge.filter(task_output)
+                            else task_output
+                        ),
+                    )
+
+                if id in output and task_output is not SKIP:
+                    return_output[id].append(task_output)
+
+            # retrieve the outgoing nodes and push them into the stack
+            for outedge, succ in successors:
+                stack.append(succ.id)
+
+        return return_output

sera/libs/directed_computing_graph/_edge.py

@@ -0,0 +1,35 @@
+from __future__ import annotations
+
+from typing import Annotated, Any, Callable, Optional
+
+from graph.interface import BaseEdge
+
+from sera.libs.directed_computing_graph._node import NodeId
+from sera.libs.directed_computing_graph._type_conversion import UnitTypeConversion
+
+
+class DCGEdge(BaseEdge[NodeId, int]):
+
+    def __init__(
+        self,
+        id: int,
+        source: NodeId,
+        target: NodeId,
+        argindex: int,
+        type_conversion: UnitTypeConversion,
+        filter_fn: Optional[Callable[[Any], bool]] = None,
+    ):
+        super().__init__(id, source, target, key=argindex)
+        self.argindex = argindex
+        self.type_conversion = type_conversion
+        self.filter_fn = filter_fn
+
+    def filter(self, value: Any) -> bool:
+        """Filter the value passing through this edge.
+
+        Returns:
+            True if the value should flow through this edge, False to block it.
+        """
+        if self.filter_fn is not None:
+            return self.filter_fn(value)
+        return True

sera/libs/directed_computing_graph/_flow.py

@@ -0,0 +1,17 @@
+from __future__ import annotations
+
+from typing import Any, Callable, Optional
+
+from sera.libs.directed_computing_graph._node import ComputeFn, ComputeFnId
+
+
+class Flow:
+    def __init__(
+        self,
+        source: list[ComputeFnId] | ComputeFnId,
+        target: ComputeFn,
+        filter_fn: Optional[Callable[[Any], bool]] = None,
+    ):
+        self.source = [source] if isinstance(source, str) else source
+        self.target = target
+        self.filter_fn = filter_fn

sera/libs/directed_computing_graph/_fn_signature.py

@@ -0,0 +1,54 @@
+from __future__ import annotations
+
+import inspect
+from dataclasses import dataclass
+from typing import Any, Callable, get_args, get_origin, get_type_hints
+
+from sera.misc import get_classpath
+
+
+@dataclass
+class FnSignature:
+    return_type: type
+    argnames: list[str]
+    argtypes: list[type]
+    default_args: dict[str, Any]  # Added this field to store default values
+    is_async: bool = False
+
+    @staticmethod
+    def parse(func: Callable) -> FnSignature:
+        sig = get_type_hints(func)
+        argnames = list(sig.keys())[:-1]
+
+        # Get the default values using inspect.signature
+        inspect_sig = inspect.signature(func)
+        defaults = {}
+        for name, param in inspect_sig.parameters.items():
+            if param.default is not inspect.Parameter.empty:
+                defaults[name] = param.default
+
+        try:
+            return FnSignature(
+                sig["return"],
+                argnames,
+                [sig[arg] for arg in argnames],
+                defaults,  # Add the default values to the signature
+                is_async=inspect.iscoroutinefunction(func),
+            )
+        except:
+            print("Cannot figure out the signature of", func)
+            print("The parsed signature is:", sig)
+            raise
+
+
+def type_to_string(_type: type) -> str:
+    """Return a fully qualified type name"""
+    origin = get_origin(_type)
+    if origin is None:
+        return get_classpath(_type)
+    return (
+        get_classpath(origin)
+        + "["
+        + ", ".join([get_classpath(arg) for arg in get_args(_type)])
+        + "]"
+    )

sera/libs/directed_computing_graph/_node.py

@@ -0,0 +1,51 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Annotated, Any, Callable
+
+from graph.interface import BaseNode
+
+from sera.libs.directed_computing_graph._fn_signature import FnSignature
+from sera.libs.directed_computing_graph._type_conversion import UnitTypeConversion
+
+
+class PartialFn:
+    def __init__(self, fn: Callable, **kwargs):
+        self.fn = fn
+        self.default_args = kwargs
+        self.signature = FnSignature.parse(fn)
+
+        argnames = set(self.signature.argnames)
+        for arg, val in self.default_args.items():
+            if arg not in argnames:
+                raise Exception(f"Argument {arg} is not in the function signature")
+            self.signature.default_args[arg] = val
+
+    def __call__(self, *args, **kwargs):
+        return self.fn(*args, **kwargs)
+
+
+ComputeFnId = Annotated[str, "ComputeFn Identifier"]
+ComputeFn = PartialFn | Callable
+NodeId = ComputeFnId
+
+
+class DCGNode(BaseNode[NodeId]):
+    id: NodeId
+    func: ComputeFn
+
+    def __init__(self, id: NodeId, func: ComputeFn):
+        super().__init__(id)
+        self.func = func
+        self.signature = self.get_signature(self.func)
+        self.type_conversions: list[UnitTypeConversion] = []
+        self.required_args: list[str] = []
+        self.required_context: list[str] = []
+        self.required_context_default_args: dict[str, Any] = {}
+
+    @staticmethod
+    def get_signature(actor: ComputeFn) -> FnSignature:
+        if isinstance(actor, PartialFn):
+            return actor.signature
+        else:
+            return FnSignature.parse(actor)

sera/libs/directed_computing_graph/_runtime.py

@@ -0,0 +1,110 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from enum import Enum
+from typing import Annotated, Any, MutableSequence, Sequence
+
+from graph.retworkx import RetworkXStrDiGraph
+
+from sera.libs.directed_computing_graph._edge import DCGEdge
+from sera.libs.directed_computing_graph._node import DCGNode, NodeId
+
+TaskKey = Annotated[tuple, "TaskKey"]
+TaskArgs = Annotated[MutableSequence, "TaskArgs"]
+
+
+class ArgValueType(Enum):
+    UNSET = "UNSET"
+    SKIP = "SKIP"
+
+
+UNSET = ArgValueType.UNSET
+SKIP = ArgValueType.SKIP
+
+
+@dataclass
+class NodeRuntime:
+    id: NodeId
+    tasks: dict[TaskKey, TaskArgs]
+    context: Sequence[Any]
+
+    graph: RetworkXStrDiGraph[int, DCGNode, DCGEdge]
+    node: DCGNode
+    indegree: int
+    # This is a mapping from parent node id to the index of the argument in the task.
+    parent2argindex: dict[str, int]
+
+    @staticmethod
+    def from_node(
+        graph: RetworkXStrDiGraph[int, DCGNode, DCGEdge],
+        node: DCGNode,
+        context: Sequence[Any],
+    ) -> NodeRuntime:
+        return NodeRuntime(
+            id=node.id,
+            tasks={},
+            context=context,
+            graph=graph,
+            node=node,
+            indegree=graph.in_degree(node.id),
+            parent2argindex={
+                edge.source: i
+                # Map parent node ID to argument index based on sorted in-edge order
+                for i, edge in enumerate(
+                    sorted(graph.in_edges(node.id), key=lambda e: e.id)
+                )
+            },
+        )
+
+    def add_task(self, key: TaskKey, args: TaskArgs) -> NodeRuntime:
+        """
+        Add a task to the node runtime.
+
+        Args:
+            key: The key identifying the task.
+            args: The arguments for the task.
+        Returns:
+            NodeRuntime: The updated node runtime with the new task added.
+        """
+        self.tasks[key] = args
+        return self
+
+    def add_task_args(
+        self, key: TaskKey, parent_node: NodeId, argvalue: Any
+    ) -> NodeRuntime:
+        """
+        Add an argument to an existing task.
+
+        Args:
+            key: The key identifying the task.
+            parent_node: Identifier of the parent node from which the argument is coming.
+            argvalue: The value of the argument to add.
+        Returns:
+            NodeRuntime: The updated node runtime with the new argument added to the task.
+        """
+        if key not in self.tasks:
+            self.tasks[key] = [UNSET] * self.indegree
+        self.tasks[key][self.parent2argindex[parent_node]] = argvalue
+        return self
+
+    def has_enough_data(self) -> bool:
+        """
+        Check if the node has enough data to execute its tasks.
+
+        Returns:
+            bool: True if the node has enough data, False otherwise.
+        """
+        return all(
+            all(arg is not UNSET for arg in args) for args in self.tasks.values()
+        )
+
+    def execute(self, task: TaskArgs) -> Any:
+        """
+        Execute a task with the given context.
+
+        Args:
+            task (TaskArgs): The arguments for the task.
+            context (dict): The context in which to execute the task.
+        """
+        norm_args = (self.node.type_conversions[i](a) for i, a in enumerate(task))
+        return self.node.func(*norm_args, *self.context)

sera/libs/directed_computing_graph/_type_conversion.py

@@ -0,0 +1,191 @@
+from __future__ import annotations
+
+import collections.abc
+import inspect
+from locale import normalize
+from types import UnionType
+from typing import (
+    Annotated,
+    Any,
+    Callable,
+    Mapping,
+    MutableMapping,
+    MutableSequence,
+    MutableSet,
+    Sequence,
+    Set,
+    TypeVar,
+    Union,
+    cast,
+    get_args,
+    get_origin,
+    get_type_hints,
+)
+
+from sera.misc import identity
+
+UnitTypeConversion = Annotated[
+    Callable[[Any], Any], "A function that convert an object of type T1 to T2"
+]
+ComposeTypeConversion = Annotated[
+    Callable[[Any, UnitTypeConversion], Any],
+    "A function that convert a generic object of type G[T1] to G[T2]",
+]
+
+
+class TypeConversion:
+    """Inspired by Rust type conversion traits. This class allows to derive a type conversion function from output of a pipe to input of another pipe."""
+
+    class UnknownConversion(Exception):
+        pass
+
+    def __init__(
+        self, type_casts: Sequence[UnitTypeConversion | ComposeTypeConversion]
+    ):
+        self.generic_single_type_conversion: dict[type, UnitTypeConversion] = {}
+        self.unit_type_conversions: dict[tuple[type, type], UnitTypeConversion] = {}
+        self.compose_type_conversion: dict[type, ComposeTypeConversion] = {}
+
+        for fn in type_casts:
+            assert not inspect.iscoroutinefunction(
+                fn
+            ), "Async conversion functions are not supported"
+            sig = get_type_hints(fn)
+            if len(sig) == 2:
+                fn = cast(UnitTypeConversion, fn)
+
+                intype = sig[[x for x in sig if x != "return"][0]]
+                outtype = sig["return"]
+
+                intype_origin = get_origin(intype)
+                intype_args = get_args(intype)
+                if (
+                    intype_origin is not None
+                    and len(intype_args) == 1
+                    and intype_args[0] is outtype
+                    and isinstance(outtype, TypeVar)
+                ):
+                    # this is a generic conversion G[T] => T
+                    self.generic_single_type_conversion[intype_origin] = fn
+                else:
+                    self.unit_type_conversions[intype, outtype] = fn
+            else:
+                assert len(sig) == 3, "Invalid type conversion function"
+                fn = cast(ComposeTypeConversion, fn)
+
+                intype = sig[[x for x in sig if x != "return"][0]]
+                outtype = sig["return"]
+                intype_origin = get_origin(intype)
+                assert intype_origin is not None
+                self.compose_type_conversion[intype_origin] = fn
+
+    def get_conversion(
+        self, source_type: type, target_type: type
+    ) -> UnitTypeConversion:
+        # handle identity conversion
+        # happen when source_type = target_type or target_type is Union[source_type, ...]
+        if source_type == target_type:
+            # source_type is target_type doesn't work with collections.abc.Sequence
+            return identity
+        if get_origin(target_type) in (Union, UnionType) and source_type in get_args(
+            target_type
+        ):
+            return identity
+
+        if (source_type, target_type) in self.unit_type_conversions:
+            # we already have a unit type conversion function for these types
+            return self.unit_type_conversions[source_type, target_type]
+
+        # check if this is a generic conversion
+        intype_origin = get_origin(source_type)
+        intype_args = get_args(source_type)
+
+        if intype_origin is None or len(intype_args) != 1:
+            raise TypeConversion.UnknownConversion(
+                f"Cannot find conversion from {source_type} to {target_type}"
+            )
+
+        outtype_origin = get_origin(target_type)
+        outtype_args = get_args(target_type)
+
+        if outtype_origin is None:
+            # we are converting G[T] => T'
+            if (
+                target_type is not intype_args[0]
+                or intype_origin not in self.generic_single_type_conversion
+            ):
+                # either T != T' or G is unkknown
+                raise TypeConversion.UnknownConversion(
+                    f"Cannot find conversion from {source_type} to {target_type}"
+                )
+            return self.generic_single_type_conversion[intype_origin]
+
+        # we are converting G[T] => G'[T']
+        if (
+            outtype_origin is not intype_origin
+            or intype_origin not in self.compose_type_conversion
+        ):
+            # either G != G' or G is unknown
+            raise TypeConversion.UnknownConversion(
+                f"Cannot find conversion from {source_type} to {target_type}"
+            )
+        # G == G' => T == T'
+        compose_func = self.compose_type_conversion[intype_origin]
+        func = self.get_conversion(intype_args[0], outtype_args[0])
+        return lambda x: compose_func(x, func)
+
+
+def patch_get_origin(t: type) -> Any:
+    """The original get_origin(typing.Sequence) returns collections.abc.Sequence.
+    Later comparing typing.Sequence[T] to collections.abc.Sequence[T] aren't equal.
+
+    This function will return typing.Sequence instead.
+    """
+    origin = get_origin(t)
+    if origin is None:
+        return origin
+    return {
+        collections.abc.Mapping: Mapping,
+        collections.abc.Sequence: Sequence,
+        collections.abc.MutableSequence: MutableSequence,
+        collections.abc.MutableMapping: MutableMapping,
+        collections.abc.Set: Set,
+        collections.abc.MutableSet: MutableSet,
+    }.get(origin, origin)
+
+
+def is_generic_type(t: type) -> bool:
+    return isinstance(t, TypeVar) or any(is_generic_type(a) for a in get_args(t))
+
+
+def align_generic_type(
+    generic_type: type, target_type: type
+) -> tuple[type, tuple[type, type]]:
+    """Return the grounded outer type, and the mapping from the TypeVar to the concrete type"""
+    if isinstance(generic_type, TypeVar):
+        return target_type, (generic_type, target_type)
+
+    origin = patch_get_origin(generic_type)
+    assert origin is not None
+    if origin != patch_get_origin(target_type):
+        raise TypeConversion.UnknownConversion(
+            f"Cannot ground generic type {generic_type} to {target_type}"
+        )
+
+    if len(get_args(generic_type)) != 1:
+        raise NotImplementedError()
+
+    gt = align_generic_type(get_args(generic_type)[0], get_args(target_type)[0])
+    return origin[gt[0]], gt[1]
+
+
+def ground_generic_type(generic_type: type, var2type: dict[TypeVar, type]) -> type:
+    if isinstance(generic_type, TypeVar):
+        return var2type[generic_type]
+
+    origin = get_origin(generic_type)
+    if origin is None:
+        # nothing to ground
+        return generic_type
+
+    return origin[*(ground_generic_type(t, var2type) for t in get_args(generic_type))]

sera/misc/__init__.py

@@ -1,8 +1,10 @@
 from sera.misc._formatter import File, Formatter
 from sera.misc._utils import (
+    LoadTableDataArgs,
     assert_isinstance,
     assert_not_null,
     filter_duplication,
+    get_classpath,
     identity,
     load_data,
     to_camel_case,
@@ -22,4 +24,6 @@ __all__ = [
     "File",
     "load_data",
     "identity",
+    "get_classpath",
+    "LoadTableDataArgs",
 ]

sera/misc/_utils.py

@@ -1,15 +1,20 @@
 from __future__ import annotations
 
 import re
+from dataclasses import dataclass
 from pathlib import Path
-from typing import Any, Callable, Iterable, Optional, TypeVar
+from typing import Any, Callable, Iterable, Optional, Sequence, Type, TypedDict, TypeVar
 
 import serde.csv
+import serde.json
 from sqlalchemy import Engine, text
 from sqlalchemy.orm import Session
 from tqdm import tqdm
 
 T = TypeVar("T")
+
+TYPE_ALIASES = {"typing.List": "list", "typing.Dict": "dict", "typing.Set": "set"}
+
 reserved_keywords = {
     "and",
     "or",
@@ -110,11 +115,19 @@ def filter_duplication(
     return new_lst
 
 
+class LoadTableDataArgs(TypedDict, total=False):
+    table: type
+    tables: Sequence[type]
+    file: Path
+    files: Sequence[Path]
+    file_deser: Callable[[Path], list[dict]]
+    record_deser: Callable[[dict], Any | list[Any]]
+
+
 def load_data(
     engine: Engine,
     create_db_and_tables: Callable[[], None],
-    table_files: list[tuple[type, Path]],
-    table_desers: dict[type, Callable[[dict], Any]],
+    table_data: Sequence[LoadTableDataArgs],
     verbose: bool = False,
 ):
     """
@@ -130,26 +143,83 @@ def load_data(
     with Session(engine) as session:
         create_db_and_tables()
 
-        for tbl, file in tqdm(table_files, disable=not verbose, desc="Loading data"):
-            if file.name.endswith(".csv"):
-                records = serde.csv.deser(file, deser_as_record=True)
+        for args in tqdm(table_data, disable=not verbose, desc="Loading data"):
+            if "table" in args:
+                tbls = [args["table"]]
+            elif "tables" in args:
+                tbls = args["tables"]
+            else:
+                raise ValueError("Either 'table' or 'tables' must be provided in args.")
+
+            if "file" in args:
+                assert isinstance(args["file"], Path), "File must be a Path object."
+                files = [args["file"]]
+            elif "files" in args:
+                assert all(
+                    isinstance(f, Path) for f in args["files"]
+                ), "Files must be Path objects."
+                files = args["files"]
+            else:
+                raise ValueError("Either 'file' or 'files' must be provided in args.")
+
+            raw_records = []
+            if "file_deser" not in args:
+                for file in files:
+                    if file.name.endswith(".csv"):
+                        raw_records.extend(serde.csv.deser(file, deser_as_record=True))
+                    elif file.name.endswith(".json"):
+                        raw_records.extend(serde.json.deser(file))
+                    else:
+                        raise ValueError(f"Unsupported file format: {file.name}")
             else:
-                raise ValueError(f"Unsupported file format: {file.name}")
-            deser = table_desers[tbl]
-            records = [deser(row) for row in records]
-            for r in tqdm(records, desc=f"load {tbl.__name__}", disable=not verbose):
-                session.merge(r)
+                for file in files:
+                    raw_records.extend(args["file_deser"](file))
+
+            deser = args["record_deser"]
+            records = [deser(row) for row in raw_records]
+            for r in tqdm(
+                records,
+                desc=f"load {', '.join(tbl.__name__ for tbl in tbls)}",
+                disable=not verbose,
+            ):
+                if isinstance(r, Sequence):
+                    for x in r:
+                        session.merge(x)
+                else:
+                    session.merge(r)
             session.flush()
 
             # Reset the sequence for each table
-            session.execute(
-                text(
-                    f"SELECT setval('{tbl.__tablename__}_id_seq', (SELECT MAX(id) FROM \"{tbl.__tablename__}\"));"
+            for tbl in tbls:
+                session.execute(
+                    text(
+                        f"SELECT setval('{tbl.__tablename__}_id_seq', (SELECT MAX(id) FROM \"{tbl.__tablename__}\"));"
+                    )
                 )
-            )
         session.commit()
 
 
 def identity(x: T) -> T:
     """Identity function that returns the input unchanged."""
     return x
+
+
+def get_classpath(type: Type | Callable) -> str:
+    if type.__module__ == "builtins":
+        return type.__qualname__
+
+    if hasattr(type, "__qualname__"):
+        return type.__module__ + "." + type.__qualname__
+
+    # typically a class from the typing module
+    if hasattr(type, "_name") and type._name is not None:
+        path = type.__module__ + "." + type._name
+        if path in TYPE_ALIASES:
+            path = TYPE_ALIASES[path]
+    elif hasattr(type, "__origin__") and hasattr(type.__origin__, "_name"):
+        # found one case which is typing.Union
+        path = type.__module__ + "." + type.__origin__._name
+    else:
+        raise NotImplementedError(type)
+
+    return path

{sera_2-1.17.0.dist-info → sera_2-1.19.0.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: sera-2
-Version: 1.17.0
+Version: 1.19.0
 Summary: 
 Author: Binh Vu
 Author-email: bvu687@gmail.com
@@ -9,13 +9,15 @@ Classifier: Programming Language :: Python :: 3
 Classifier: Programming Language :: Python :: 3.12
 Classifier: Programming Language :: Python :: 3.13
 Requires-Dist: black (==25.1.0)
-Requires-Dist: codegen-2 (>=2.11.1,<3.0.0)
+Requires-Dist: codegen-2 (>=2.12.0,<3.0.0)
+Requires-Dist: graph-wrapper (>=1.7.2,<2.0.0)
 Requires-Dist: isort (==6.0.1)
 Requires-Dist: litestar (>=2.15.1,<3.0.0)
 Requires-Dist: loguru (>=0.7.0,<0.8.0)
 Requires-Dist: msgspec (>=0.19.0,<0.20.0)
-Requires-Dist: serde2 (>=1.9.0,<2.0.0)
+Requires-Dist: serde2 (>=1.9.2,<2.0.0)
 Requires-Dist: sqlalchemy[asyncio] (>=2.0.41,<3.0.0)
+Requires-Dist: tqdm (>=4.67.1,<5.0.0)
 Requires-Dist: typer (>=0.12.3,<0.13.0)
 Project-URL: Repository, https://github.com/binh-vu/sera
 Description-Content-Type: text/markdown

{sera_2-1.17.0.dist-info → sera_2-1.19.0.dist-info}/RECORD

@@ -8,8 +8,14 @@ sera/libs/api_helper.py,sha256=47y1kcwk3Xd2ZEMnUj_0OwCuUmgwOs5kYrE95BDVUn4,5411
 sera/libs/api_test_helper.py,sha256=3tRr8sLN4dBSrHgKAHMmyoENI0xh7K_JLel8AvujU7k,1323
 sera/libs/base_orm.py,sha256=5hOH_diUeaABm3cpE2-9u50VRqG1QW2osPQnvVHIhIA,3365
 sera/libs/base_service.py,sha256=AX1WoTHte6Z_birkkfagkNE6BrCLTlTjQE4jEsKEaAY,5152
-sera/libs/dag/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-sera/libs/dag/_dag.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+sera/libs/directed_computing_graph/__init__.py,sha256=xiF5_I1y9HtQ-cyq02iwkRYgEZvxBB8YIvysCHCLBco,1290
+sera/libs/directed_computing_graph/_dcg.py,sha256=AGTzKVSl-EsSOJlNKPOA1Io7pIxfq0SMXuumq1IExl0,14902
+sera/libs/directed_computing_graph/_edge.py,sha256=iBq6cpLWWyuD99QWTHVEh8naWUJrR4WJJuq5iuCrwHo,1026
+sera/libs/directed_computing_graph/_flow.py,sha256=6v39yKPIDYrQ3KvFqjeAWs88-oQSnDTaED2F3LF2z_I,478
+sera/libs/directed_computing_graph/_fn_signature.py,sha256=73iPUITcRKW0-l6sqjwMSk_FZnJESaKOmUKDGHTOh9Q,1598
+sera/libs/directed_computing_graph/_node.py,sha256=9FsKceW_hq6RYaC7d5YKF5aSXmbAcj-LGakh_GCNgHw,1597
+sera/libs/directed_computing_graph/_runtime.py,sha256=76Ccl1Rj31SkzRJPWFvYNu9ZzUABoeHp5v3tfScekcI,3319
+sera/libs/directed_computing_graph/_type_conversion.py,sha256=_XGvDidOJVmHS4gqdPlhJGzdV34YtNiPF5Kr2nV6ZgE,6806
 sera/libs/middlewares/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 sera/libs/middlewares/auth.py,sha256=r6aix1ZBwxMd1Jv5hMCTB8a_gFOJQ6egvxIrf3DWEOs,2323
 sera/libs/middlewares/uscp.py,sha256=H5umW8iEQSCdb_MJ5Im49kxg1E7TpxSg1p2_2A5zI1U,2600
@@ -20,9 +26,9 @@ sera/make/make_python_api.py,sha256=iXGbKQ3IJvsY1ur_fhurr_THFNnH66E3Wl85o0emUbw,
 sera/make/make_python_model.py,sha256=Nc4vDGgM8icgWBqzNnMgEkLadf5EsZwbbHs3WLW9_co,62778
 sera/make/make_python_services.py,sha256=0ZpWLwQ7Nwfn8BXAikAB4JRpNknpSJyJgY5b1cjtxV4,2073
 sera/make/make_typescript_model.py,sha256=1ouYFCeqOlwEzsGBiXUn4VZtLJjJW7GSacdOSlQzhjI,67012
-sera/misc/__init__.py,sha256=mPKkik00j3tO_m45VPDJBjm8K85NpymRPl36Kh4hBn8,473
+sera/misc/__init__.py,sha256=rOmGMv7QNzpSKZSyxChbRmEnBr3O443UlLGS0FIs3AI,561
 sera/misc/_formatter.py,sha256=aCGYL08l8f3aLODHxSocxBBwkRYEo3K1QzCDEn3suj0,1685
-sera/misc/_utils.py,sha256=pGYv8p7m7opiDTLYbsPrhF0YA4WjFff7beMQQZ9NnEs,4095
+sera/misc/_utils.py,sha256=O0Qh_BZXX7kmzhFBTnnxJT4BUxQbwNovhepSwhh33Ow,6532
 sera/models/__init__.py,sha256=vJC5Kzo_N7wd16ocNPy1VvAZDGNiWeiAhWJ4ihATKvA,780
 sera/models/_class.py,sha256=1J4Bd_LanzhhDWwZFHWGtFYD7lupe_alaB3D02ebNDI,2862
 sera/models/_collection.py,sha256=ZnQEriKC4X88Zz48Kn1AVZKH-1_l8OgWa-zf2kcQOOE,1414
@@ -36,6 +42,6 @@ sera/models/_parse.py,sha256=ciTLzCkO0q6xA1R_rHbnYJYK3Duo2oh56WeuwxXwJaI,12392
 sera/models/_property.py,sha256=9yMDxrmbyuF6-29lQjiq163Xzwbk75TlmGBpu0NLpkI,7485
 sera/models/_schema.py,sha256=VxJEiqgVvbXgcSUK4UW6JnRcggk4nsooVSE6MyXmfNY,1636
 sera/typing.py,sha256=o_DKfSvs8JpNRQ0kdaTc3BbfdkvibY3uY4tJRt-n2fQ,1023
-sera_2-1.17.0.dist-info/METADATA,sha256=aIaXid2dkyX8P9nty-1eFHFBuH0Cpy34vOGDi1wTFkI,852
-sera_2-1.17.0.dist-info/WHEEL,sha256=b4K_helf-jlQoXBBETfwnf4B04YC67LOev0jo4fX5m8,88
-sera_2-1.17.0.dist-info/RECORD,,
+sera_2-1.19.0.dist-info/METADATA,sha256=22HLRJ7LXZYcBduEokfXaTnePGNCUIyARqVDZFP404k,936
+sera_2-1.19.0.dist-info/WHEEL,sha256=b4K_helf-jlQoXBBETfwnf4B04YC67LOev0jo4fX5m8,88
+sera_2-1.19.0.dist-info/RECORD,,