dstklib 1.0.0__py3-none-any.whl

dstk/workflow_tools.py

@@ -0,0 +1,257 @@
+from functools import wraps
+import warnings
+import inspect
+from copy import deepcopy
+
+from .lib_types.dstk_types import Function, MethodSpec
+from inspect import Signature, BoundArguments
+from typing import Any, cast, Callable, Type, TypeAlias, TypeGuard
+
+class WorkflowManager:
+    """
+    Manages the execution of processing methods in workflow mode.
+
+    Tracks workflow state, controls stage transitions, and stores intermediate
+    results for chained method execution with enforced sequencing and unit context.
+    """
+
+    def __init__(self) -> None:
+        """
+        Initializes WorkflowManager with given attributes.
+        """
+
+        self._flow: bool
+        self._current_stage: str
+        self._processing_unit: str
+
+        self._stages: list[str]
+
+        self._start: Any
+        self._end: Any
+
+        self._called_methods: list[str] = []
+
+    def _set_workflow(self, input_arg: Any) -> None:
+        """
+        Initializes workflow mode based on the presence of input arguments.
+
+        Sets workflow state and starting point if all required inputs are provided.
+        Issues a warning and disables workflow mode if inputs are partially missing.
+
+        :param input_args: A dictionary of input argument names and their values. If all values are non-None, workflow mode is activated.
+        :param input_source: The initial data source to store when starting the workflow.
+        """
+
+        if input_arg is not None:
+            self._start = input_arg
+            self._current_stage = "start"
+            self._flow = True
+        else:
+            self._flow = False
+
+    @property
+    def result(self) -> Any:
+        """
+        Returns the current output of the processing workflow.
+
+        Use this property to retrieve the final result after a chain of workflow method calls. It safely copies the internal state (if possible) to prevent side effects.
+
+        :return: The result of the most recent workflow stage.
+        """
+        
+        result: Any = getattr(self, f"_{self._current_stage}")
+        try:
+            copy: Any = deepcopy(result)
+            return copy
+        except:
+            return result
+
+
+class WorkflowBuilder: 
+    """
+    Automates the execution of a sequence of methods on a WorkflowManager subclass.
+
+    :param work_class: A subclass of WorkflowManager representing the workflow to execute.
+    :param method_representation: A dictionary mapping method names to their keyword arguments.
+    :param result: If True, returns the result of the workflow. Else, returns the instance of the working class. Defaults to True.
+    """
+
+    def __init__(self, work_class: Type[WorkflowManager], method_representation: MethodSpec, result: bool = True):
+        """
+        Initializes WorkflowBuilder with given attributes.
+        """
+
+        self.work_class: type = work_class
+        self.methods: MethodSpec = method_representation
+        self.result: bool = result
+
+    def __call__(self, *args, **kwargs) -> Any:
+        workflow: Type[WorkflowManager] = self.work_class(*args, **kwargs)
+
+        for key, value in self.methods.items():
+            method: Callable = getattr(workflow, key)
+            method(**value)
+        
+        return workflow.result if self.result else workflow
+        
+
+def workflow(input_arg: str, input_process: str, output_process: str,  input_attrs: dict[str, Any] | None = None, next_stage: str | None = None, set_unit: str | None = None,) -> Callable[[Function], Function]:
+    """
+    Enables workflow execution for a method by automatically injecting inputs,
+    storing outputs, and transitioning stages when in workflow mode.
+
+    :param input_arg: Name of the keyword argument to inject into the method.
+    :param input_process: Attribute name to retrieve the input data from if not provided.
+    :param output_process: Attribute name to store the method's output in workflow mode.
+    :param input_attrs: Optional mapping of argument names to extract from input data. Supports str for attribute access or nested dicts for deep lookup. Defaults to None.
+    :param next_stage: Optional name of the next workflow stage to transition to after method execution. Defaults to None.
+    :param set_unit: Optional name of the processing unit to set for the next workflow step. Defaults to None.
+
+    :return: A decorator that wraps the method with workflow logic.
+    :raises ValueError: If the user is not in workflow mode and he did not passed thte input arg. Also, if the value of input_attrs is different from None, str or dict.
+    """
+
+    def decorator(method: Function) -> Function:
+        @wraps(method)
+        def wrapper(self, *args, **kwargs) -> Any:
+            method_name: str = method.__name__
+
+            if input_arg not in kwargs:
+                if not self._flow:
+                    raise ValueError(f"{input_arg} must be provided if not using workflow mode")
+
+                input_data: Any = getattr(self, input_process)
+
+                if input_attrs:
+                    for key, value in input_attrs.items():
+                        if value is None:
+                            kwargs[key] = input_data
+                        elif isinstance(value, str):
+                            kwargs[key] = getattr(input_data, value)
+                        elif isinstance(value, dict):
+                            for attr_name, subattr in value.items():
+                                mapping = getattr(input_data, attr_name)[subattr] 
+                                kwargs[key] = mapping
+                        else:
+                            raise ValueError(f"Type {type(value)} of value bot recognized")
+                else:
+                    kwargs[input_arg] = input_data
+
+            result: Any = method(self, *args, **kwargs)
+
+            if self._flow:
+                setattr(self, output_process, result)
+                if next_stage:
+                    self._current_stage = next_stage
+                if set_unit:
+                    self._processing_unit = set_unit
+                if next_stage == "end":
+                  warnings.warn(UserWarning(f"After calling method {method_name} you must necessarily call result to continue with analysis. Further chaining will result in error."))
+
+                self._called_methods.append(method_name)
+                return self
+
+            return result
+
+        return cast(Function, wrapper)
+
+    return decorator
+
+def requires(stages: list[str], unit: str | None = None, multiple_calls: bool = False) -> Callable[[Function], Function]:
+    """
+    Ensures a method is only callable in workflow mode at allowed stages and units,
+    and prevents repeated calls to the same method.
+
+    :param stages: A list of stages where the method is allowed to be used.
+    :param unit: The required processing unit; if specified, the method can only run when the current unit matches. Defaults to None.
+
+    :return: A decorator that enforces workflow constraints on the wrapped method.
+    """
+
+    def decorator(method: Function) -> Function:
+        @wraps(method)
+        def wrapper(self, *args, **kwargs) -> Any:
+            method_name: str = method.__name__
+
+            if self._flow:
+                if method_name in self._called_methods and not multiple_calls:
+                    raise RuntimeError(f"Method {method_name} already called. You can only call each method exactly once in workflow mode.")
+                if not hasattr(self, "_current_stage"):
+                    raise RuntimeError("Current phase is not initialized.")
+                if self._current_stage not in stages:
+                    raise RuntimeError(
+                        f"Method '{method_name}' requires stages '{', '.join(stages)}' but current phase is '{self._current_stage}'"
+                    )
+                if unit and unit != self._processing_unit:
+                    raise RuntimeError(
+                        f"Method '{method_name}' can only be used when the you are processing by {unit}, but you are currently processing by {self._processing_unit}."
+                    )
+                
+            return method(self, *args, **kwargs)
+
+        return cast(Function, wrapper)
+
+    return decorator
+
+def accepts_generic(*, type_checker: Callable, input_arg: str, accepts: bool, intercept: bool, interceptor: Callable, input_type: TypeAlias, custom_error_message: str = "") -> Callable[[Function], Function]:
+    """
+    A generic decorator factory that conditionally intercepts and transforms a method's input based on runtime type checks.
+
+    This utility enables creating flexible, reusable decorators (like `accepts_sentences` or `accepts_tags`) by delegating the interception logic to a custom `interceptor` and input validation to a `type_checker`.
+
+    If the input value (specified by `input_arg`) passes the `type_checker`, and both `accepts` and `intercept` are True, the `interceptor` is called instead of the original method. Otherwise, the original method is called directly. If the input is of the expected type but `accepts` is False, a `ValueError` is raised.
+
+    :param type_checker: Function that checks whether the input value is of the expected structure/type.
+    :param input_arg: The name of the argument to inspect in the decorated method.
+    :param accepts: Whether the method is allowed to handle this type of input.
+    :param intercept: Whether the input should be transformed before calling the method.
+    :param interceptor: Function that handles the interception logic, replacing the original method call when triggered.
+    :param input_type: A human-readable type description used in error messages and type casts. This is for documentation/static typing purposes only.
+    :param custom_error_message: Optional custom message to append to error if input is rejected.
+
+    :returns: Callable: A decorator that wraps the target method with conditional input handling logic.
+
+    :raises: ValueError: If the input matches the expected type but `accepts` is False.
+    """
+
+    def decorator(method: Function) -> Function:
+        @wraps(method)
+        def wrapper(self, *args, **kwargs) -> Any:
+            signature: Signature = inspect.signature(method)
+            bound_args: BoundArguments = signature.bind(self, *args, **kwargs)
+            bound_args.apply_defaults()
+
+            input_value: Any = bound_args.arguments.get(input_arg, None)
+
+            if type_checker(input_value) and (accepts and intercept):
+                filtered_kwargs: dict[str, Any] = kwargs.copy()
+                filtered_kwargs.pop(input_arg, None)
+
+                return interceptor(self, input_value=input_value, method=method, *args, **filtered_kwargs)
+            elif not type_checker(input_value) or (accepts and not intercept):
+                return cast(input_type, method(self, *args, **kwargs))
+            else:
+                raise ValueError(f"Method {method.__name__} does not accept {input_type} as input. {custom_error_message}")
+                
+        return cast(Function, wrapper)
+    return decorator
+
+def is_method_spec(spec: Any) -> TypeGuard[MethodSpec]:
+    """
+    If spec is of type MethodSpec, returns True. Else, returns False.
+
+    :param spec: A dict to check its type.
+    """
+
+    return (
+        isinstance(spec, dict) and
+        all(
+            isinstance(method_name, str) and
+            isinstance(params, dict) and 
+            all(
+                isinstance(key, str)
+                for key in params.keys()
+            )
+            for method_name, params in spec.items()
+        )
+    )