sigima 0.0.1.dev0__py3-none-any.whl → 1.0.0__py3-none-any.whl

sigima/proc/decorator.py

@@ -0,0 +1,330 @@
+# Copyright (c) DataLab Platform Developers, BSD 3-Clause license, see LICENSE file.
+
+"""
+.. Computation function decorator and utilities
+(see parent package :mod:`sigima.computation`)
+"""
+
+from __future__ import annotations
+
+import dataclasses
+import functools
+import importlib
+import inspect
+import os.path as osp
+import pkgutil
+import sys
+import typing
+from typing import Callable, Literal, TypeVar
+
+import guidata.dataset as gds
+import makefun
+
+if sys.version_info >= (3, 10):
+    # Use ParamSpec from typing module in Python 3.10+
+    from typing import ParamSpec
+else:
+    # Use ParamSpec from typing_extensions module in Python < 3.10
+    from typing_extensions import ParamSpec
+
+# NOTE: Parameter classes should NOT be included in __all__ to avoid Sphinx
+# cross-reference conflicts. All parameter classes are re-exported through
+# sigima.params module which serves as the single source of truth for the
+# public API. Only utility functions should be exported from this module.
+__all__ = [
+    "computation_function",
+    "find_computation_functions",
+    "get_computation_metadata",
+    "is_computation_function",
+]
+
+# Marker attribute used by @computation_function and introspection
+COMPUTATION_METADATA_ATTR = "__computation_function_metadata"
+
+P = ParamSpec("P")
+R = TypeVar("R")
+
+
+@dataclasses.dataclass(frozen=True)
+class ComputationMetadata:
+    """Metadata for a computation function.
+
+    Attributes:
+        name: The name of the computation function.
+        description: A description or docstring for the computation function.
+    """
+
+    name: str
+    description: str
+
+
+def _make_computation_wrapper(
+    f: Callable,
+    ds_cls: type,
+    ds_param: inspect.Parameter,
+    params: list,
+    ds_items: list,
+    new_sig: inspect.Signature,
+    signature_info: str,
+    metadata: ComputationMetadata,
+) -> Callable:
+    """
+    Create a computation function wrapper supporting both DataSet and expanded-kwarg
+    signatures.
+
+    Args:
+        f: The original function.
+        ds_cls: The DataSet class type.
+        ds_param: The DataSet parameter in the signature.
+        params: The full function signature parameters.
+        ds_items: The DataSet's items (parameters).
+        new_sig: The explicit signature (with kwargs) to expose.
+        signature_info: The Sphinx docstring note to append.
+        metadata: ComputationMetadata to attach to the wrapper.
+
+    Returns:
+        The wrapped function.
+    """
+
+    @makefun.with_signature(new_sig)
+    @functools.wraps(f)
+    def wrapper(*args, **kwargs):
+        """
+        Dispatch function supporting both DataSet parameter and expanded keyword
+        arguments.
+
+        Behavior:
+            - If a DataSet object is provided, it is always used and keyword arguments
+              for DataSet items are ignored.
+            - If no DataSet is provided, DataSet items are constructed from keyword
+              arguments.
+
+        Returns:
+            Result of the original computation function.
+        """
+        ba = new_sig.bind(*args, **kwargs)
+        ba.apply_defaults()
+        ds_obj = ba.arguments.get(ds_param.name, None)
+        ds_item_names = set(item.get_name() for item in ds_items)
+
+        if isinstance(ds_obj, ds_cls):
+            # DataSet object provided: ignore any keyword arguments for its items
+            pass
+        else:
+            # DataSet instance not provided: build from keyword arguments
+            ds_kwargs = {
+                k: ba.arguments.pop(k)
+                for k in list(ba.arguments.keys())
+                if k in ds_item_names
+            }
+            ds_obj = ds_cls.create(**ds_kwargs)
+
+        # Build the final positional argument list for the original function
+        final_args = []
+        for p in params:
+            if p is ds_param:
+                final_args.append(ds_obj)
+            else:
+                final_args.append(ba.arguments.get(p.name, None))
+        return f(*final_args)
+
+    # Attach dynamic Sphinx docstring and signature
+    doc = f.__doc__ or ""
+    if not doc.endswith("\n"):
+        doc += "\n"
+    wrapper.__doc__ = doc + signature_info
+    wrapper.__signature__ = new_sig
+    setattr(wrapper, COMPUTATION_METADATA_ATTR, metadata)
+    return wrapper
+
+
+def computation_function(
+    *,
+    name: typing.Optional[str] = None,
+    description: typing.Optional[str] = None,
+) -> Callable[[Callable[P, R]], Callable[P, R]]:
+    """
+    Decorator to mark a function as a Sigima computation function.
+
+    This decorator enables two calling conventions:
+      1. With a guidata DataSet object as a parameter (classic style).
+      2. With the DataSet items passed as individual keyword arguments (expanded style).
+
+    The decorator ensures:
+      - An explicit and informative function signature (including all DataSet items as
+        keyword arguments).
+      - A Sphinx-friendly docstring documenting both call styles.
+      - Pickle-compatibility (crucial for multiprocessing).
+      - Conflict detection if both DataSet instance and expanded keyword arguments are
+        used simultaneously.
+
+    Args:
+        name: Optional custom name for metadata.
+        description: Optional custom description or docstring.
+
+    Returns:
+        The decorated, enhanced computation function.
+    """
+
+    def decorator(f: Callable[P, R]) -> Callable[P, R]:
+        # Gather signature and typing information
+        sig = inspect.signature(f)
+        params = list(sig.parameters.values())
+        try:
+            type_hints = typing.get_type_hints(f)
+        except Exception:  # pylint: disable=broad-except
+            type_hints = {}
+
+        # Find DataSet parameter if any
+        ds_param = None
+        ds_cls = None
+        for p in params:
+            annot = type_hints.get(p.name, p.annotation)
+            if (
+                annot is not inspect.Signature.empty
+                and isinstance(annot, type)
+                and issubclass(annot, gds.DataSet)
+                and annot.__name__ not in ("SignalObj", "ImageObj")
+            ):
+                ds_param = p
+                ds_cls = annot
+                break
+
+        # If a DataSet param is present, expand signature and docstring
+        if ds_cls is not None:
+            # Build signature exposing all DataSet items as keyword-only parameters
+            ds_items: list[gds.DataItem] = ds_cls._items  # pylint: disable=W0212
+            item_names = [item.get_name() for item in ds_items]
+            items = []
+            for item in ds_items:
+                if item.get_name() not in [p.name for p in params]:
+                    # Support ChoiceItem as Literal if available
+                    if hasattr(gds, "ChoiceItem") and isinstance(item, gds.ChoiceItem):
+                        choice_data = item.get_prop("data", "choices")
+                        choices = [v[0] for v in choice_data]
+                        item_type = Literal[tuple(choices)]
+                    else:
+                        item_type = item.type
+                    items.append(
+                        inspect.Parameter(
+                            item.get_name(),
+                            inspect.Parameter.KEYWORD_ONLY,
+                            annotation=item_type,
+                            default=item.get_default(),
+                        )
+                    )
+            # DataSet parameter remains positional-or-keyword, but optional
+            # (default=None)
+            base_params = []
+            for p in params:
+                if p is ds_param:
+                    base_params.append(
+                        inspect.Parameter(
+                            p.name,
+                            inspect.Parameter.POSITIONAL_OR_KEYWORD,
+                            annotation=p.annotation,
+                            default=None,
+                        )
+                    )
+                else:
+                    base_params.append(p)
+            new_params = base_params + items
+            new_sig = sig.replace(parameters=new_params)
+            param_class_name = ds_cls.__name__
+            kwarg_example = ", ".join(f"{name}=..." for name in item_names)
+            # Sphinx-style docstring describing both call conventions
+            signature_info = (
+                f".. note::\n\n"
+                f"   This computation function can be called in two ways:\n\n"
+                f"   1. With a parameter ``{param_class_name}`` object:\n\n"
+                f"      .. code-block:: python\n\n"
+                f"         param = {param_class_name}.create({kwarg_example})\n"
+                f"         func(obj, param)\n\n"
+                f"   2. Or, with keyword arguments directly:\n\n"
+                f"      .. code-block:: python\n\n"
+                f"         func(obj, {kwarg_example})\n\n"
+                f"   Both styles are fully supported and equivalent.\n\n"
+            )
+            metadata = ComputationMetadata(
+                name=name or f.__name__,
+                description=description or f.__doc__,
+            )
+            return _make_computation_wrapper(
+                f, ds_cls, ds_param, params, ds_items, new_sig, signature_info, metadata
+            )
+
+        # No DataSet parameter: simple passthrough
+        @functools.wraps(f)
+        def wrapper(*args, **kwargs):
+            return f(*args, **kwargs)
+
+        metadata = ComputationMetadata(
+            name=name or f.__name__,
+            description=description or f.__doc__,
+        )
+        setattr(wrapper, COMPUTATION_METADATA_ATTR, metadata)
+        return wrapper
+
+    return decorator
+
+
+def is_computation_function(function: Callable) -> bool:
+    """Check if a function is a Sigima computation function.
+
+    Args:
+        function: The function to check.
+
+    Returns:
+        True if the function is a Sigima computation function, False otherwise.
+    """
+    return getattr(function, COMPUTATION_METADATA_ATTR, None) is not None
+
+
+def get_computation_metadata(function: Callable) -> ComputationMetadata:
+    """Get the metadata of a Sigima computation function.
+
+    Args:
+        function: The function to get metadata from.
+
+    Returns:
+        Computation function metadata.
+
+    Raises:
+        ValueError: If the function is not a Sigima computation function.
+    """
+    metadata = getattr(function, COMPUTATION_METADATA_ATTR, None)
+    if not isinstance(metadata, ComputationMetadata):
+        raise ValueError(
+            f"The function {function.__name__} is not a Sigima computation function."
+        )
+    return metadata
+
+
+def find_computation_functions() -> list[tuple[str, Callable]]:
+    """Find all computation functions in the `sigima.proc` package.
+
+    This function uses introspection to locate all functions decorated with
+    `@computation_function` in the `sigima.proc` package and its subpackages.
+
+    Args:
+        module: Optional module to search in. If None, the current module is used.
+
+    Returns:
+        A list of tuples, each containing the function name and the function object.
+    """
+    functions = []
+    objs = []
+    for _, modname, _ in pkgutil.walk_packages(
+        path=[osp.dirname(__file__)], prefix=".".join(__name__.split(".")[:-1]) + "."
+    ):
+        try:
+            module = importlib.import_module(modname)
+        except Exception:  # pylint: disable=broad-except
+            continue
+        for name, obj in inspect.getmembers(module, inspect.isfunction):
+            if is_computation_function(obj):
+                if obj in objs:  # Avoid double entries for the same function
+                    continue
+                objs.append(obj)
+                functions.append((modname, name, obj.__doc__))
+    return functions