stepfunction 0.1.0__py3-none-any.whl → 0.2.0__py3-none-any.whl

stepfunction/constants/visualizer.py

@@ -46,3 +46,8 @@ DEFAULT_VISUALIZER_SUB_STEP_FUNCTION_CLASS_STYLE = (
     "fill:#f5f5f5,stroke:#333,stroke-dasharray: 5 5"
 )
 """str: The Mermaid classDef style applied to sub-step function nodes."""
+
+# Node ID sanitization
+
+VISUALIZER_INVALID_NODE_ID_CHARS = r"[^0-9A-Za-z_]"
+"""str: Regex matching characters not allowed in a Mermaid node ID; matches are replaced with "_"."""

stepfunction/core/serializer/__init__.py

@@ -0,0 +1,3 @@
+from .serializer import decode_step_function, encode_step_function
+
+__all__ = ["decode_step_function", "encode_step_function"]

stepfunction/core/serializer/serializer.py

@@ -0,0 +1,170 @@
+"""Encode/decode logic for declarative (JSON-able dict) StepFunction specs.
+
+Author: Vineeth Penugonda
+"""
+
+from typing import Any, Callable, Dict, Optional, cast
+
+from stepfunction.core.step_function.step_function import StepFunction
+from stepfunction.exceptions.step_errors import (
+    UnregisteredFunctionError,
+    UnserializableStepError,
+)
+from stepfunction.registry.step_registry import StepRegistry
+from stepfunction.registry.step_registry import registry as default_registry
+from stepfunction.types.step_types import StepParams
+
+
+def _name_for_or_raise(
+    step_registry: StepRegistry, func: Callable[[Any], Any], step_name: str
+) -> str:
+    name = step_registry.name_for(func)
+    if name is None:
+        raise UnregisteredFunctionError(
+            f"The function used in step '{step_name}' is not registered in the "
+            "given registry, so its name can't be determined for export. "
+            "Register it with stepfunction.registry.step_registry.register_step() first."
+        )
+    return name
+
+
+def _encode_step(
+    step_name: str, step: StepParams, step_registry: StepRegistry
+) -> Dict[str, Any]:
+    if step["step_type"] is not None:
+        raise UnserializableStepError(step_name, step["step_type"])
+
+    encoded: Dict[str, Any] = {
+        "next_step": step["next_step"],
+        "on_failure": step["on_failure"],
+        "parallel": step["parallel"],
+        "stop_on_failure": step["stop_on_failure"],
+    }
+
+    if step["is_sub_step_function"]:
+        sub_step_function = cast(StepFunction, step["sub_step_function"])
+        encoded["sub_step_function"] = encode_step_function(
+            sub_step_function, step_registry
+        )
+        return encoded
+
+    if step["parallel"]:
+        func_map = cast(Dict[str, Callable[[Any], Any]], step["func"])
+        encoded["func"] = {
+            slot: _name_for_or_raise(step_registry, fn, step_name)
+            for slot, fn in func_map.items()
+        }
+    else:
+        func = cast(Callable[[Any], Any], step["func"])
+        encoded["func"] = _name_for_or_raise(step_registry, func, step_name)
+
+    branch = step["branch"]
+    if branch is not None:
+        if callable(branch):
+            encoded["branch"] = _name_for_or_raise(step_registry, branch, step_name)
+        else:
+            encoded["branch"] = {str(key): value for key, value in branch.items()}
+
+    return encoded
+
+
+def encode_step_function(
+    sf: StepFunction, step_registry: Optional[StepRegistry] = None
+) -> Dict[str, Any]:
+    """Export ``sf`` as a JSON-able dict.
+
+    Recurses into nested sub-step-functions by calling itself again, so
+    arbitrary nesting depth is handled without special-casing.
+
+    Raises:
+        UnserializableStepError: If any step was built from a BaseStep
+            instance (RetryStep, TimeoutStep, WaitStep, or a custom
+            BaseStep subclass) — not yet supported.
+        UnregisteredFunctionError: If a step or branch function used in
+            ``sf`` has no registered name in ``step_registry``.
+    """
+    step_registry = step_registry or default_registry
+
+    return {
+        "name": sf.name,
+        "start_step": sf.current_step,
+        "steps": {
+            step_name: _encode_step(step_name, step, step_registry)
+            for step_name, step in sf.steps.items()
+        },
+    }
+
+
+def decode_step_function(
+    data: Dict[str, Any],
+    step_registry: Optional[StepRegistry] = None,
+    _validate: bool = True,
+) -> StepFunction:
+    """Reconstruct a StepFunction from a dict produced by ``encode_step_function``.
+
+    Rebuilds the workflow purely through ``add_step``/``add_sub_step_function``/
+    ``set_start_step`` — the same public API a user would call by hand — and
+    recurses into nested "sub_step_function" entries by calling itself again.
+    Validates exactly once, at the very end, at the outermost level only:
+    ``StepFunction.validate()`` already recurses into sub-step functions and
+    reports a readable breadcrumb across nesting levels, so a malformed
+    nested spec still fails fast without the decoder needing its own
+    recursive validation pass.
+
+    Raises:
+        UnregisteredFunctionError: If a referenced function name isn't
+            registered in ``step_registry`` (defaults to the package's
+            default singleton registry if not given).
+        ValueError: If the reconstructed workflow fails validate().
+    """
+    step_registry = step_registry or default_registry
+
+    sf = StepFunction(data["name"])
+
+    for step_name, step_data in data["steps"].items():
+        if "sub_step_function" in step_data:
+            sub_step_function = decode_step_function(
+                step_data["sub_step_function"], step_registry, _validate=False
+            )
+            sf.add_sub_step_function(
+                step_name,
+                sub_step_function=sub_step_function,
+                next_step=step_data.get("next_step"),
+                on_failure=step_data.get("on_failure"),
+            )
+            continue
+
+        func_spec = step_data["func"]
+        if isinstance(func_spec, dict):
+            func: Any = {
+                slot: step_registry.get(ref) for slot, ref in func_spec.items()
+            }
+        else:
+            func = step_registry.get(func_spec)
+
+        branch_spec = step_data.get("branch")
+        if branch_spec is None:
+            branch: Any = None
+        elif isinstance(branch_spec, dict):
+            branch = dict(branch_spec)
+        else:
+            branch = step_registry.get(branch_spec)
+
+        sf.add_step(
+            step_name,
+            func,
+            next_step=step_data.get("next_step"),
+            on_failure=step_data.get("on_failure"),
+            branch=branch,
+            parallel=step_data.get("parallel", False),
+            stop_on_failure=step_data.get("stop_on_failure", False),
+        )
+
+    start_step = data.get("start_step")
+    if start_step is not None:
+        sf.set_start_step(start_step)
+
+    if _validate:
+        sf.validate()
+
+    return sf

stepfunction/core/step_function/step_function.py

@@ -3,18 +3,24 @@
 Author: Vineeth Penugonda
 """
 
-from asyncio import gather, get_running_loop
-from inspect import iscoroutinefunction
-from typing import Any, Callable, Dict, List, Optional, Union, cast
+from asyncio import CancelledError, gather, get_running_loop
+from inspect import isawaitable, iscoroutinefunction
+from time import monotonic
+from typing import TYPE_CHECKING, Any, Callable, Dict, List, Optional, Union, cast
 
 from stepfunction.constants.enums import StepFunctionStatus
 from stepfunction.exceptions.step_errors import (
     ParallelStepExecutionError,
     StepExecutionError,
 )
+from stepfunction.hooks import StepEvent, StepFunctionHooks, WorkflowEvent
 from stepfunction.steps.base import BaseStep
 from stepfunction.types.step_types import StepParams
 from stepfunction.utils.logger import setup_logger
+from stepfunction.utils.utils import utc_now
+
+if TYPE_CHECKING:
+    from stepfunction.registry.step_registry import StepRegistry
 
 
 class StepFunction:
@@ -33,6 +39,7 @@ class StepFunction:
 
     Properties:
         name (str): The name of the step function.
+        hooks (Optional[StepFunctionHooks]): The lifecycle hooks set on this step function, if any.
         steps (Dict[str, StepParams]): A dictionary containing the steps of the workflow.
         last_result (Any): The result of the last step.
         context (Dict[str, Any]): Stores step names and results, including exceptions if any occur.
@@ -49,6 +56,9 @@ class StepFunction:
         set_start_step(name):
             Set the start step of the workflow.
 
+        set_hooks(hooks):
+            Set the lifecycle hooks (a StepFunctionHooks) notified as steps start, succeed and fail.
+
         add_sub_step_function(name, sub_step_function, next_step=None, on_failure=None):
             Add a sub-step function to be executed as a step.
 
@@ -66,6 +76,14 @@ class StepFunction:
         visualize_to_string():
             Return a string representation of the workflow for visualization.
 
+        to_dict(step_registry=None) / to_json(indent=2, step_registry=None):
+            Export the workflow as a JSON-able dict / JSON string, resolving step and
+            branch functions to their registered name via stepfunction.registry.step_registry.
+
+        from_dict(data, step_registry=None) / from_json(json_str, step_registry=None):
+            Classmethods that reconstruct a StepFunction from a dict/JSON string
+            produced by to_dict()/to_json(), resolving function names via the registry.
+
     Protected Methods:
         _execute_step(func, input_value):
             Execute a single step, handling asynchronous and synchronous functions appropriately.
@@ -108,11 +126,25 @@ class StepFunction:
 
     Status Example:
         status = step_function.status  # Will be StepFunctionStatus.INITIALIZED, StepFunctionStatus.RUNNING, StepFunctionStatus.COMPLETED, or StepFunctionStatus.FAILED.
+
+    Hooks Example:
+        class RunRecorder(StepFunctionHooks):
+            async def on_step_success(self, event):
+                print(event.step, event.duration, event.next_step)
+
+        step_function = StepFunction("MyStepFunction", hooks=RunRecorder())
     """
 
-    def __init__(self, name: str):
+    def __init__(self, name: str, hooks: Optional[StepFunctionHooks] = None):
         self.__name = name  # Name of the step function
 
+        self.__hooks = hooks  # Lifecycle hooks set on this step function
+
+        # Hooks and workflow path for the current run; a sub-step function
+        # without hooks of its own inherits its parent's
+        self.__run_hooks: Optional[StepFunctionHooks] = None
+        self.__run_path: List[str] = [name]
+
         self.__steps: Dict[str, StepParams] = {}  # Steps of the workflow
         self.__current_step = None  # The current step being executed
 
@@ -181,7 +213,13 @@ class StepFunction:
             raise ValueError(f"Step '{name}' already exists in steps")
 
         async def sub_func(last_result):
-            await sub_step_function.execute(initial_input=last_result)
+            await sub_step_function._execute(
+                initial_input=last_result,
+                hooks=sub_step_function.hooks
+                if sub_step_function.hooks is not None
+                else self.__run_hooks,
+                path=self.__run_path + [_sub_path_label(name, sub_step_function)],
+            )
             self.__context.update(sub_step_function.context)
 
             if sub_step_function.status == StepFunctionStatus.FAILED:
@@ -208,6 +246,10 @@ class StepFunction:
 
         self.__current_step = name
 
+    def set_hooks(self, hooks: Optional[StepFunctionHooks]):
+        """Set the lifecycle hooks notified as the workflow runs (None removes them)."""
+        self.__hooks = hooks
+
     def validate(self):
         """Validate the workflow configuration before execution."""
         self._validate(path=[self.__name])
@@ -246,16 +288,33 @@ class StepFunction:
             if step["is_sub_step_function"]:
                 sub_step_function = cast("StepFunction", step["sub_step_function"])
                 sub_step_function._validate(
-                    path=path + [f"{step_name} ({sub_step_function.name})"]
-                    if step_name != sub_step_function.name
-                    else path + [step_name]
+                    path=path + [_sub_path_label(step_name, sub_step_function)]
                 )
 
     async def execute(self, initial_input: Any = None):
         """Execute the workflow."""
 
+        await self._execute(initial_input, hooks=self.__hooks, path=[self.__name])
+
+    async def _execute(
+        self,
+        initial_input: Any,
+        hooks: Optional[StepFunctionHooks],
+        path: List[str],
+    ):
+        """Execute the workflow, reporting to ``hooks`` under the workflow ``path``.
+
+        ``path`` is the chain of step-function names from the outermost
+        workflow down to this one, as in _validate(); sub-step functions are
+        run through here so their events carry the full path.
+        """
+
         self.validate()
 
+        self.__run_hooks = hooks
+        self.__run_path = path
+        workflow = " -> ".join(path)
+
         self.__status = StepFunctionStatus.RUNNING
 
         self.__logger.debug(
@@ -264,96 +323,174 @@ class StepFunction:
 
         self.__last_result = initial_input
 
-        while self.__current_step:
-            step = self.__steps[self.__current_step]
-            try:
-                if step["parallel"]:
-                    results = await self._execute_parallel(
-                        cast(Dict[str, Callable[[Any], Any]], step["func"]),
-                        step["stop_on_failure"],
-                    )
-
-                    self.__last_result = results
-                    self.__context[self.__current_step] = results
+        run_started_at = utc_now()
+        run_started = monotonic()
+        await self._emit(
+            "on_workflow_start",
+            WorkflowEvent(
+                workflow=workflow,
+                status=self.__status,
+                input=initial_input,
+                started_at=run_started_at,
+            ),
+        )
 
-                    self.__logger.info(
-                        f"Parallel step '{self.__current_step}' succeeded with results: {results}"
-                    )
-                else:
-                    result = await self._execute_step(
-                        cast(Callable[[Any], Any], step["func"]), self.__last_result
-                    )
+        run_error: Optional[BaseException] = None
+        try:
+            while self.__current_step:
+                await self._execute_current_step(workflow)
+        except BaseException as exc:
+            run_error = exc
+            self.__status = StepFunctionStatus.FAILED
+            raise
+        finally:
+            await self._emit(
+                "on_workflow_end",
+                WorkflowEvent(
+                    workflow=workflow,
+                    status=self.__status,
+                    input=initial_input,
+                    started_at=run_started_at,
+                    output=self.__last_result,
+                    error=run_error,
+                    finished_at=utc_now(),
+                    duration=monotonic() - run_started,
+                ),
+            )
 
-                    self.__last_result = result
-                    self.__context[self.__current_step] = result
+    async def _execute_current_step(self, workflow: str):
+        """Execute the current step, route to the next one and report both to the hooks."""
+
+        step_name = cast(str, self.__current_step)
+        step = self.__steps[step_name]
+        step_input = self.__last_result
+
+        started_at = utc_now()
+        started = monotonic()
+        await self._emit(
+            "on_step_start",
+            StepEvent(
+                workflow=workflow,
+                step=step_name,
+                task=None,
+                input=step_input,
+                started_at=started_at,
+            ),
+        )
 
-                    self.__logger.info(f"Step '{self.__current_step}' succeeded")
+        def _ended(**fields: Any) -> StepEvent:
+            return StepEvent(
+                workflow=workflow,
+                step=step_name,
+                task=None,
+                input=step_input,
+                started_at=started_at,
+                finished_at=utc_now(),
+                duration=monotonic() - started,
+                **fields,
+            )
 
-                next_step = None
+        try:
+            if step["parallel"]:
+                results = await self._execute_parallel(
+                    cast(Dict[str, Callable[[Any], Any]], step["func"]),
+                    step["stop_on_failure"],
+                    workflow=workflow,
+                    step_name=step_name,
+                )
 
-                if step["branch"]:
-                    if callable(step["branch"]):
-                        next_step = step["branch"](self.__last_result)
-                    else:
-                        next_step = step["branch"].get(self.__last_result)
+                self.__last_result = results
+                self.__context[step_name] = results
 
-                    if next_step is None:
-                        self.__logger.warning(
-                            f"Step '{self.__current_step}': branch did not resolve to a next step "
-                            f"for result '{self.__last_result}'. Workflow will end."
-                        )
-                    elif next_step not in self.__steps:
-                        self.__logger.warning(
-                            f"Step '{self.__current_step}': branch resolved to '{next_step}' "
-                            "which does not exist in steps. This will cause a failure."
-                        )
+                self.__logger.info(
+                    f"Parallel step '{step_name}' succeeded with results: {results}"
+                )
+            else:
+                result = await self._execute_step(
+                    cast(Callable[[Any], Any], step["func"]), self.__last_result
+                )
 
-                self.__current_step = next_step or step["next_step"]
+                self.__last_result = result
+                self.__context[step_name] = result
 
-            except Exception as exc:
-                self.__logger.exception(
-                    f"Step '{self.__current_step}' failed. Exception: {exc}"
-                )
+                self.__logger.info(f"Step '{step_name}' succeeded")
 
-                exc_value = exc.args[0] if exc.args else exc
+            next_step = None
 
-                self.__context[cast(str, self.__current_step)] = exc_value
+            if step["branch"]:
+                if callable(step["branch"]):
+                    next_step = step["branch"](self.__last_result)
+                else:
+                    next_step = step["branch"].get(self.__last_result)
 
-                if step["on_failure"]:
-                    self.__logger.exception(
-                        f"Executing failure step: {step['on_failure']} for '{self.__current_step}'"
+                if next_step is None:
+                    self.__logger.warning(
+                        f"Step '{step_name}': branch did not resolve to a next step "
+                        f"for result '{self.__last_result}'. Workflow will end."
+                    )
+                elif next_step not in self.__steps:
+                    self.__logger.warning(
+                        f"Step '{step_name}': branch resolved to '{next_step}' "
+                        "which does not exist in steps. This will cause a failure."
                     )
 
-                    self.__current_step = step["on_failure"]
-                    self.__last_result = exc_value
+            self.__current_step = next_step or step["next_step"]
 
-                    self.__status = StepFunctionStatus.FAILED
+        except CancelledError as exc:
+            self.__logger.warning(f"Step '{step_name}' was cancelled")
 
-                    self.__logger.debug(
-                        f"StepFunction - {self.__name} - Status - {self.__status.value}"
-                    )
-                else:
-                    self.__logger.exception(
-                        f"No failure step defined for '{self.__current_step}'. Raising Exception."
-                    )
+            self.__status = StepFunctionStatus.FAILED
 
-                    self.__status = StepFunctionStatus.FAILED
+            await self._emit("on_step_failure", _ended(error=exc))
 
-                    self.__logger.debug(
-                        f"StepFunction - {self.__name} - Status - {self.__status.value}"
-                    )
+            raise
+
+        except Exception as exc:
+            self.__logger.exception(f"Step '{step_name}' failed. Exception: {exc}")
+
+            exc_value = exc.args[0] if exc.args else exc
+
+            self.__context[step_name] = exc_value
+
+            self.__status = StepFunctionStatus.FAILED
 
-                    raise StepExecutionError(exc)
+            self.__logger.debug(
+                f"StepFunction - {self.__name} - Status - {self.__status.value}"
+            )
 
-            if not self.__current_step and self.__status != StepFunctionStatus.FAILED:
-                # If no more steps, mark as COMPLETED
+            await self._emit(
+                "on_step_failure", _ended(error=exc, next_step=step["on_failure"])
+            )
 
-                self.__status = StepFunctionStatus.COMPLETED
+            if step["on_failure"]:
+                self.__logger.exception(
+                    f"Executing failure step: {step['on_failure']} for '{step_name}'"
+                )
 
-                self.__logger.debug(
-                    f"StepFunction - {self.__name} - Status - {self.__status.value}"
+                self.__current_step = step["on_failure"]
+                self.__last_result = exc_value
+            else:
+                self.__logger.exception(
+                    f"No failure step defined for '{step_name}'. Raising Exception."
                 )
 
+                raise StepExecutionError(exc)
+
+        else:
+            await self._emit(
+                "on_step_success",
+                _ended(output=self.__last_result, next_step=self.__current_step),
+            )
+
+        if not self.__current_step and self.__status != StepFunctionStatus.FAILED:
+            # If no more steps, mark as COMPLETED
+
+            self.__status = StepFunctionStatus.COMPLETED
+
+            self.__logger.debug(
+                f"StepFunction - {self.__name} - Status - {self.__status.value}"
+            )
+
     async def _execute_step(self, func: Callable, input_value: Any):
         """Execute a single step, handling async functions."""
         if iscoroutinefunction(func):
@@ -362,30 +499,72 @@ class StepFunction:
             return func(input_value)
 
     async def _execute_parallel(
-        self, func_dict: Dict[str, Callable[[Any], Any]], stop_on_failure: bool = False
+        self,
+        func_dict: Dict[str, Callable[[Any], Any]],
+        stop_on_failure: bool = False,
+        workflow: Optional[str] = None,
+        step_name: Optional[str] = None,
     ):
-        """Execute the steps in parallel without blocking the event loop."""
+        """Execute the steps in parallel without blocking the event loop.
+
+        Each task is reported to the hooks as its own step event, with
+        ``task`` set to its name.
+        """
         loop = get_running_loop()
         results = {}
         errors = []
+        step_input = self.__last_result
+
+        async def _run_one(task: str, func: Callable[[Any], Any]) -> Any:
+            started_at = utc_now()
+            started = monotonic()
+            await self._emit(
+                "on_step_start",
+                StepEvent(
+                    workflow=cast(str, workflow),
+                    step=cast(str, step_name),
+                    task=task,
+                    input=step_input,
+                    started_at=started_at,
+                ),
+            )
 
-        async def _run_one(func: Callable[[Any], Any]) -> Any:
-            if iscoroutinefunction(func):
-                return await func(self.__last_result)
-            return await loop.run_in_executor(None, func, self.__last_result)
+            def _ended(**fields: Any) -> StepEvent:
+                return StepEvent(
+                    workflow=cast(str, workflow),
+                    step=cast(str, step_name),
+                    task=task,
+                    input=step_input,
+                    started_at=started_at,
+                    finished_at=utc_now(),
+                    duration=monotonic() - started,
+                    **fields,
+                )
+
+            try:
+                if iscoroutinefunction(func):
+                    result = await func(step_input)
+                else:
+                    result = await loop.run_in_executor(None, func, step_input)
+            except BaseException as exc:
+                await self._emit("on_step_failure", _ended(error=exc))
+                raise
+
+            await self._emit("on_step_success", _ended(output=result))
+            return result
 
         task_results = await gather(
-            *[_run_one(func) for func in func_dict.values()],
+            *[_run_one(task, func) for task, func in func_dict.items()],
             return_exceptions=True,
         )
 
-        for step_name, result in zip(func_dict.keys(), task_results):
+        for task, result in zip(func_dict.keys(), task_results):
             if isinstance(result, Exception):
-                self.__logger.exception(f"Parallel task '{step_name}' failed: {result}")
-                results[step_name] = result.args[0] if result.args else result
-                errors.append((step_name, result))
+                self.__logger.exception(f"Parallel task '{task}' failed: {result}")
+                results[task] = result.args[0] if result.args else result
+                errors.append((task, result))
             else:
-                results[step_name] = result
+                results[task] = result
 
         if errors:
             self.__logger.error(f"Some parallel tasks failed: {errors}")
@@ -396,6 +575,25 @@ class StepFunction:
 
         return results
 
+    async def _emit(self, hook_name: str, event: Union[StepEvent, WorkflowEvent]):
+        """Call a hook for this run, awaiting it if it's a coroutine.
+
+        A hook that raises is logged and ignored: hooks observe the workflow,
+        they never change its routing or status.
+        """
+        hooks = self.__run_hooks
+        if hooks is None:
+            return
+
+        try:
+            result = getattr(hooks, hook_name)(event)
+            if isawaitable(result):
+                await result
+        except Exception:
+            self.__logger.exception(
+                f"StepFunction - {self.__name} - Hook '{hook_name}' raised; ignoring"
+            )
+
     def visualize(self):
         """Visualize the workflow."""
         from stepfunction.core.visualizer import Visualizer
@@ -430,11 +628,80 @@ class StepFunction:
 
         return visualizer.render_step_function_to_string()
 
+    def to_dict(self, step_registry: Optional["StepRegistry"] = None) -> Dict[str, Any]:
+        """Export this workflow as a JSON-able dict.
+
+        Function references (step funcs, parallel-step funcs, and branch
+        callables) are encoded as their registered string name — see
+        ``stepfunction.registry.step_registry``. Nested sub-step-functions are
+        encoded recursively.
+
+        Call this before execute(): if the workflow has already run,
+        "start_step" reflects wherever the execution cursor ended up
+        (often None after a completed run), not the original start step.
+
+        Raises:
+            UnserializableStepError: If any step was built from a BaseStep
+                instance (RetryStep, TimeoutStep, WaitStep, or a custom
+                BaseStep subclass) — not yet supported.
+            UnregisteredFunctionError: If a step or branch function used in
+                this workflow has no registered name in ``step_registry``.
+        """
+        from stepfunction.core.serializer import encode_step_function
+
+        return encode_step_function(self, step_registry)
+
+    def to_json(
+        self, indent: Optional[int] = 2, step_registry: Optional["StepRegistry"] = None
+    ) -> str:
+        """Export this workflow as a JSON string. See to_dict() for details
+        and caveats."""
+        from json import dumps
+
+        return dumps(self.to_dict(step_registry), indent=indent)
+
+    @classmethod
+    def from_dict(
+        cls, data: Dict[str, Any], step_registry: Optional["StepRegistry"] = None
+    ) -> "StepFunction":
+        """Reconstruct a StepFunction from a dict produced by to_dict().
+
+        Function name references in ``data`` are resolved against
+        ``step_registry`` (defaults to the package's default registry,
+        ``stepfunction.registry.step_registry.registry``, if not given). Validates the
+        reconstructed workflow before returning, so a malformed spec fails
+        fast with a readable breadcrumb error rather than only failing
+        later at execute() time.
+
+        Raises:
+            UnregisteredFunctionError: If a referenced function name isn't
+                registered.
+            ValueError: If the reconstructed workflow fails validate().
+        """
+        from stepfunction.core.serializer import decode_step_function
+
+        return decode_step_function(data, step_registry)
+
+    @classmethod
+    def from_json(
+        cls, json_str: str, step_registry: Optional["StepRegistry"] = None
+    ) -> "StepFunction":
+        """Reconstruct a StepFunction from a JSON string produced by to_json().
+        See from_dict() for details."""
+        from json import loads
+
+        return cls.from_dict(loads(json_str), step_registry)
+
     @property
     def name(self):
         """Returns the name of the step function."""
         return self.__name
 
+    @property
+    def hooks(self) -> Optional[StepFunctionHooks]:
+        """Returns the lifecycle hooks set on the step function, if any."""
+        return self.__hooks
+
     @property
     def steps(self):
         """Returns the steps of the step function."""
@@ -455,8 +722,22 @@ class StepFunction:
         """Returns the status of the step function."""
         return self.__status
 
+    @property
+    def current_step(self):
+        """Returns the step set by set_start_step(), advanced by execute()
+        as the workflow progresses."""
+        return self.__current_step
+
     def __str__(self):
         return f"StepFunction(name={self.__name}, NoOfSteps={len(self.__steps)}, CurrentStep={self.__current_step}, Status={self.__status})"
 
     def __repr__(self):
         return str(self)
+
+
+def _sub_path_label(step_name: str, sub_step_function: StepFunction) -> str:
+    """How a sub-step function appears in a workflow path: its step name, plus
+    its own name when the two differ (e.g. "SubStep (SUB_FLOW)")."""
+    if step_name != sub_step_function.name:
+        return f"{step_name} ({sub_step_function.name})"
+    return step_name

stepfunction/core/visualizer/visualizer.py

@@ -36,11 +36,10 @@ from stepfunction.constants.visualizer import (
     DEFAULT_VISUALIZER_SUB_STEP_FUNCTION_CLASS,
     DEFAULT_VISUALIZER_SUB_STEP_FUNCTION_CLASS_STYLE,
     DEFAULT_VISUALIZER_SUCCESS_EDGE_LABEL,
+    VISUALIZER_INVALID_NODE_ID_CHARS,
 )
 from stepfunction.types.step_types import StepParams
 
-_INVALID_NODE_ID_CHARS = r"[^0-9A-Za-z_]"
-
 
 def _quote(label: str) -> str:
     """Escape and quote a string for use as a Mermaid edge label."""
@@ -261,7 +260,7 @@ class Visualizer:
         if step_name in self.__node_ids:
             return self.__node_ids[step_name]
 
-        sanitized = re_sub(_INVALID_NODE_ID_CHARS, "_", step_name) or "step"
+        sanitized = re_sub(VISUALIZER_INVALID_NODE_ID_CHARS, "_", step_name) or "step"
         if sanitized[0].isdigit():
             sanitized = f"_{sanitized}"
 

stepfunction/exceptions/step_errors.py

@@ -12,3 +12,39 @@ class ParallelStepExecutionError(Exception):
     def __init__(self, exc: list[tuple[str, Exception]]):
         self.message = f"Parallel step generated an exception: {exc}"
         super().__init__(self.message)
+
+
+class DuplicateRegistrationError(Exception):
+    """Error raised when a function is registered under a name that is already taken."""
+
+    def __init__(self, name: str):
+        self.message = f"A function is already registered under the name '{name}'."
+        super().__init__(self.message)
+
+
+class UnregisteredFunctionError(Exception):
+    """Error raised when a step spec references a function that has no
+    registered name, in either direction: decoding a name that was never
+    registered, or encoding a function that was never given a name."""
+
+    def __init__(self, message: str):
+        self.message = message
+        super().__init__(self.message)
+
+
+class UnserializableStepError(Exception):
+    """Error raised when a step cannot be represented in a declarative spec.
+
+    Currently applies to any step built from a BaseStep instance (e.g.
+    RetryStep, TimeoutStep, WaitStep) — only the built closure survives
+    add_step(), not the original instance/config needed to reconstruct it.
+    """
+
+    def __init__(self, step_name: str, step_type: object):
+        self.message = (
+            f"Step '{step_name}' cannot be serialized: it was built from a "
+            f"BaseStep instance (step_type={step_type}). Serializing "
+            "BaseStep-derived steps (RetryStep, TimeoutStep, WaitStep, or "
+            "custom BaseStep subclasses) is not yet supported."
+        )
+        super().__init__(self.message)

stepfunction/hooks/__init__.py

@@ -0,0 +1,4 @@
+from .events import StepEvent, WorkflowEvent
+from .step_function_hooks import StepFunctionHooks
+
+__all__ = ["StepEvent", "StepFunctionHooks", "WorkflowEvent"]

stepfunction/hooks/events.py

@@ -0,0 +1,73 @@
+"""Events passed to lifecycle hooks.
+
+Author: Vineeth Penugonda
+"""
+
+from dataclasses import dataclass
+from datetime import datetime
+from typing import Any, Optional
+
+from stepfunction.constants.enums import StepFunctionStatus
+
+
+@dataclass(frozen=True)
+class StepEvent:
+    """Describes one step (or one task of a parallel step) as it starts or ends.
+
+    Attributes:
+        workflow (str): Path of the step function running the step, from the
+            outermost workflow down, e.g. "FLOW -> SubStep (SUB_FLOW)".
+        step (str): Name of the step.
+        task (Optional[str]): Name of the task within a parallel step; None for
+            ordinary steps and for the parallel step itself.
+        input (Any): The value the step was called with.
+        started_at (datetime): When the step started (UTC).
+        output (Any): The step's result. Set on success only.
+        error (Optional[BaseException]): The exception the step raised. Set on
+            failure only — the original exception, not the value stored in context.
+        next_step (Optional[str]): Where the workflow goes next: the branch or
+            next_step target on success, the on_failure target on failure, or
+            None if the workflow ends (or raises) here.
+        finished_at (Optional[datetime]): When the step ended (UTC). None on start.
+        duration (Optional[float]): Seconds the step took, from a monotonic
+            clock. None on start.
+    """
+
+    workflow: str
+    step: str
+    task: Optional[str]
+    input: Any
+    started_at: datetime
+    output: Any = None
+    error: Optional[BaseException] = None
+    next_step: Optional[str] = None
+    finished_at: Optional[datetime] = None
+    duration: Optional[float] = None
+
+
+@dataclass(frozen=True)
+class WorkflowEvent:
+    """Describes a step function run as it starts or ends.
+
+    Attributes:
+        workflow (str): Path of the step function, from the outermost workflow down.
+        status (StepFunctionStatus): RUNNING on start; COMPLETED or FAILED on end.
+        input (Any): The initial input the run was called with.
+        started_at (datetime): When the run started (UTC).
+        output (Any): The last step's result. Set on end only.
+        error (Optional[BaseException]): The exception execute() raised (e.g. a
+            StepExecutionError wrapping the step's error — the step's own failure
+            event has the original). None when the run completed, or failed into
+            an on_failure step.
+        finished_at (Optional[datetime]): When the run ended (UTC). None on start.
+        duration (Optional[float]): Seconds the run took. None on start.
+    """
+
+    workflow: str
+    status: StepFunctionStatus
+    input: Any
+    started_at: datetime
+    output: Any = None
+    error: Optional[BaseException] = None
+    finished_at: Optional[datetime] = None
+    duration: Optional[float] = None

stepfunction/hooks/step_function_hooks.py

@@ -0,0 +1,51 @@
+"""Base class for observing a step function's lifecycle.
+
+Author: Vineeth Penugonda
+"""
+
+from stepfunction.hooks.events import StepEvent, WorkflowEvent
+
+
+class StepFunctionHooks:
+    """Receives events as a step function runs.
+
+    Subclass it and override the methods you need; the rest do nothing. Each
+    method may be a regular function or a coroutine function. Pass an instance
+    to ``StepFunction(name, hooks=...)`` or ``set_hooks()``.
+
+    Hooks observe; they don't steer. An exception raised by a hook is logged
+    and ignored, so it never changes the workflow's routing or status. A
+    sub-step function without hooks of its own uses its parent's.
+
+    Order of events for one run:
+        on_workflow_start
+        for each step: on_step_start, then on_step_success or on_step_failure
+            (a parallel step also reports each task between its own start and end,
+            with ``event.task`` set)
+        on_workflow_end
+
+    If the run is cancelled, the running step and the workflow end with
+    ``error`` set to the ``asyncio.CancelledError``, which is then re-raised.
+
+    Example:
+        class PrintHooks(StepFunctionHooks):
+            async def on_step_success(self, event: StepEvent):
+                print(f"{event.step} took {event.duration:.3f}s -> {event.next_step}")
+
+        sf = StepFunction("MyStepFunction", hooks=PrintHooks())
+    """
+
+    def on_workflow_start(self, event: WorkflowEvent):
+        """Called once the workflow is validated, before its first step."""
+
+    def on_workflow_end(self, event: WorkflowEvent):
+        """Called when the workflow completes, fails, raises or is cancelled."""
+
+    def on_step_start(self, event: StepEvent):
+        """Called before a step (or a parallel task) runs."""
+
+    def on_step_success(self, event: StepEvent):
+        """Called after a step (or a parallel task) returns and its next step is resolved."""
+
+    def on_step_failure(self, event: StepEvent):
+        """Called after a step (or a parallel task) raises."""

stepfunction/registry/step_registry.py

@@ -0,0 +1,97 @@
+"""Function registry for declarative (JSON) workflow loading.
+
+Author: Vineeth Penugonda
+"""
+
+from typing import Any, Callable, Dict, Optional
+
+from stepfunction.exceptions.step_errors import (
+    DuplicateRegistrationError,
+    UnregisteredFunctionError,
+)
+
+
+class StepRegistry:
+    """A name -> callable registry used to resolve string function references
+    in a declarative (JSON) workflow spec back into real Python callables.
+
+    Step functions and branch-router functions share one namespace, since
+    both are simply ``Callable[[Any], Any]`` — the calling context
+    (``add_step`` vs. ``branch=``), not the function itself, determines the
+    role a registered callable plays in a given workflow.
+
+    A module-level default instance (``registry``) covers the common case of
+    a single process-wide registry. Construct your own ``StepRegistry()`` for
+    isolated namespaces (e.g. per-tenant, or test isolation).
+    """
+
+    def __init__(self):
+        self.__functions: Dict[str, Callable[[Any], Any]] = {}
+
+    def register(
+        self, name: str, func: Optional[Callable[[Any], Any]] = None
+    ) -> Callable[[Any], Any]:
+        """Register ``func`` under ``name``.
+
+        Usable as a direct call (``registry.register("my_step", my_step)``)
+        or as a decorator (``@registry.register("my_step")``).
+
+        Raises:
+            DuplicateRegistrationError: If ``name`` is already registered.
+        """
+        if func is not None:
+            self.__add(name, func)
+            return func
+
+        def decorator(fn: Callable[[Any], Any]) -> Callable[[Any], Any]:
+            self.__add(name, fn)
+            return fn
+
+        return decorator
+
+    def __add(self, name: str, func: Callable[[Any], Any]) -> None:
+        if name in self.__functions:
+            raise DuplicateRegistrationError(name)
+        self.__functions[name] = func
+
+    def unregister(self, name: str) -> None:
+        """Remove a registration if present. No error if ``name`` isn't registered."""
+        self.__functions.pop(name, None)
+
+    def get(self, name: str) -> Callable[[Any], Any]:
+        """Look up a registered callable by name.
+
+        Raises:
+            UnregisteredFunctionError: If ``name`` has not been registered.
+        """
+        try:
+            return self.__functions[name]
+        except KeyError:
+            raise UnregisteredFunctionError(
+                f"No function is registered under the name '{name}'. Register it "
+                "with stepfunction.registry.step_registry.register_step() (or your "
+                "own StepRegistry instance) before loading this workflow spec."
+            ) from None
+
+    def name_for(self, func: Callable[[Any], Any]) -> Optional[str]:
+        """Reverse lookup: the registered name for ``func``, or ``None`` if it
+        isn't registered under this registry."""
+        for registered_name, registered_func in self.__functions.items():
+            if registered_func is func:
+                return registered_name
+        return None
+
+    def __contains__(self, name: str) -> bool:
+        return name in self.__functions
+
+    def clear(self) -> None:
+        """Remove all registrations. Primarily useful for test isolation."""
+        self.__functions.clear()
+
+
+# Module-level default singleton, covering the common case of one
+# process-wide registry.
+registry = StepRegistry()
+
+register_step = registry.register
+get_step = registry.get

stepfunction/utils/utils.py

@@ -1,5 +1,6 @@
 """This module contains utility functions for the stepfunction package."""
 
+from datetime import datetime, timezone
 from os import getenv
 from typing import Optional
 
@@ -16,3 +17,8 @@ def get_environment_variable(name: str, default: Optional[str] = None) -> Option
         str: The value of the environment variable.
     """
     return getenv(name, default)
+
+
+def utc_now() -> datetime:
+    """Returns the current time as a timezone-aware UTC datetime."""
+    return datetime.now(timezone.utc)

{stepfunction-0.1.0.dist-info → stepfunction-0.2.0.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: stepfunction
-Version: 0.1.0
+Version: 0.2.0
 Summary: Step Function Workflow Orchestration Library
 Author: Vineeth Penugonda
 License-Expression: MIT
@@ -33,6 +33,39 @@ Dynamic: license-file
 - **Branching Logic**: Direct workflows based on conditions.
 - **Sub-Step Functions**: Modularize workflows by embedding sub-step functions.
 - **Visualization**: Integrated support for visualizing the workflow graph as a Mermaid flowchart.
+- **Lifecycle Hooks**: Get notified as the workflow and each step start, succeed or fail — with inputs, outputs, errors, routing and timings — to record run history, metrics or traces.
+
+## Lifecycle Hooks
+
+Subclass `StepFunctionHooks`, override the events you need (sync or async), and pass an instance to the step function:
+
+```python
+from stepfunction.core.step_function import StepFunction
+from stepfunction.hooks import StepEvent, StepFunctionHooks
+
+
+class RunRecorder(StepFunctionHooks):
+    async def on_step_start(self, event: StepEvent):
+        print(f"{event.workflow}: {event.step} started with {event.input!r}")
+
+    async def on_step_success(self, event: StepEvent):
+        print(f"{event.step} -> {event.output!r} in {event.duration:.3f}s, next: {event.next_step}")
+
+    async def on_step_failure(self, event: StepEvent):
+        print(f"{event.step} failed: {event.error!r}, next: {event.next_step}")
+
+
+sf = StepFunction("MyStepFunction", hooks=RunRecorder())
+```
+
+Events: `on_workflow_start`, `on_step_start`, `on_step_success`, `on_step_failure`, `on_workflow_end`.
+
+- `StepEvent` carries the workflow path, step name, parallel `task` name (if any), `input`, `output`, the original `error`, the resolved `next_step`, UTC `started_at` / `finished_at` and `duration` in seconds.
+- `WorkflowEvent` carries the workflow path, `status`, initial `input`, final `output`, the `error` that ended the run (if `execute()` raised) and timings.
+- Each task of a parallel step is reported as its own step event, with `task` set.
+- A sub-step function without hooks of its own reports to its parent's hooks; its events carry the nested path, e.g. `"FLOW -> SubStep (SUB_FLOW)"`.
+- Hooks observe; they don't steer. An exception raised by a hook is logged and ignored.
+- If the run is cancelled, the running step and the workflow are reported with the `CancelledError`, which is then re-raised.
 
 ## Installation
 

{stepfunction-0.1.0.dist-info → stepfunction-0.2.0.dist-info}/RECORD

@@ -1,12 +1,20 @@
+stepfunction/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 stepfunction/constants/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 stepfunction/constants/enums.py,sha256=K4_xKGa0NApw6Ai4VkfGW1Sc5fEqwFPpuyezy6YexzY,240
-stepfunction/constants/visualizer.py,sha256=impswU0JpqgpzTu4_a26FVLK4tII4ptbUaopFb2ZObM,1815
+stepfunction/constants/visualizer.py,sha256=MHi7MafiXGr6wb2ug7-czRG-CiM-xO5wniN6EFZBRzA,1994
+stepfunction/core/serializer/__init__.py,sha256=PvxFmVDK8EnFBqZXgYBELO1uobT5vkFFntRzlGx7cJo,127
+stepfunction/core/serializer/serializer.py,sha256=k5Uu881Amoy0Kv2mPwiv7kcNOnydtfOEQeBRu3aloj4,6006
 stepfunction/core/step_function/__init__.py,sha256=4L5ufP13A7hXBDBYrD45pfEsV9bsa6zbSkSnHpi-tyM,68
-stepfunction/core/step_function/step_function.py,sha256=qPfHSEP1LFLR6_hRVxxD6duALaDcTkIRWO_HrdYAg7s,17666
+stepfunction/core/step_function/step_function.py,sha256=whoKdcjayIHnu8wSEnB06ufhhXLnQNrmvb_7I2qXcT4,27809
 stepfunction/core/visualizer/__init__.py,sha256=pgNI6LPnY9Pm9e_YhMn-5CVmUzp8mvoytbDZ_b3McSw,61
-stepfunction/core/visualizer/visualizer.py,sha256=XM-kSoIPha-rYM2T-ZYYoo6-ttcnodn-832MX9Gg6d0,12892
+stepfunction/core/visualizer/visualizer.py,sha256=hhgZt5y98KiKLFtxhc_W7U0NwQ48F6WmAmN_N2_0GKc,12897
 stepfunction/exceptions/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-stepfunction/exceptions/step_errors.py,sha256=yNJt6tiO8LkiSxeHoSZBpaKA7qkqdsXPnOd3AMt3C8A,486
+stepfunction/exceptions/step_errors.py,sha256=jN71QXMtuekNy9YDMrwkcOxY9RR9Wssbk-zKeGVW0I4,1923
+stepfunction/hooks/__init__.py,sha256=omNFwha6-07ggGZHnDTZCbZNEM26qvffIxtRvdGp8RU,159
+stepfunction/hooks/events.py,sha256=gIqaEJqQRro7VPyqfDEgdLQNK627iX3E6fOYz2Cf9lg,2886
+stepfunction/hooks/step_function_hooks.py,sha256=JK6wLb_uns3DoECt6tQJGCiSNMhbgZ71HEcNADq4pig,2011
+stepfunction/registry/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+stepfunction/registry/step_registry.py,sha256=znpqXkgJQnOwCg7K3h6BmHdZrrDIFhsnyOYzdW8iVQE,3430
 stepfunction/steps/__init__.py,sha256=DvzYMJiMY-5O2tWhCa82LSbqIa80VXfUuLgFYyAOdi8,256
 stepfunction/steps/retry_step.py,sha256=Y5cIWi4MqVSeXUVxHKXW0TkYPBl7bOcGGL1cxH_26wY,2855
 stepfunction/steps/timeout_step.py,sha256=0_wwAAZcwC61fKSkT4TWRef-mUOT3A2T_-Y1pTXuhd8,2379
@@ -20,9 +28,9 @@ stepfunction/types/step_types.py,sha256=pMA8xY_a3yNmpDul8BBqa-0_Fpp3QNIMFpxAgn_t
 stepfunction/utils/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 stepfunction/utils/constants.py,sha256=VZefE-OhNXOyg28wPu62VZPvReublyqWS515-TbYWo0,391
 stepfunction/utils/logger.py,sha256=0L8F8oO5HiYNqcuV8G8UD4cTdCpmObQYrkvwvrF1S-8,1789
-stepfunction/utils/utils.py,sha256=Zw9-sKE2LOghxyiKImqGWSv4nHZ9nI8yEptp1xaYfHk,563
-stepfunction-0.1.0.dist-info/licenses/LICENSE,sha256=sLQZUABreUtT56S4w592eAmWiFONtLKZG-KZWf-_dvg,1074
-stepfunction-0.1.0.dist-info/METADATA,sha256=QG6O4Dkz66awYt9zyBhTBKwYbBFWmQQ4VWBuQyrr3rw,2836
-stepfunction-0.1.0.dist-info/WHEEL,sha256=aeYiig01lYGDzBgS8HxWXOg3uV61G9ijOsup-k9o1sk,91
-stepfunction-0.1.0.dist-info/top_level.txt,sha256=XqLJSUeYSmzTNMiPj-IESxp1om-LdZjefT4Un2orQD4,13
-stepfunction-0.1.0.dist-info/RECORD,,
+stepfunction/utils/utils.py,sha256=CrJuiWQrVs3re569wlJFODkmZ2fZcQaUS7Vxm56188c,739
+stepfunction-0.2.0.dist-info/licenses/LICENSE,sha256=sLQZUABreUtT56S4w592eAmWiFONtLKZG-KZWf-_dvg,1074
+stepfunction-0.2.0.dist-info/METADATA,sha256=ojVNltqxMd6afHuj6RpORW05ldaJ-AB6SI_fO59Kc_8,4737
+stepfunction-0.2.0.dist-info/WHEEL,sha256=YVMoNqKzERt-wjUZwJ33xBGAwnFl-4cqbYkTtWa4itE,91
+stepfunction-0.2.0.dist-info/top_level.txt,sha256=XqLJSUeYSmzTNMiPj-IESxp1om-LdZjefT4Un2orQD4,13
+stepfunction-0.2.0.dist-info/RECORD,,

{stepfunction-0.1.0.dist-info → stepfunction-0.2.0.dist-info}/WHEEL

@@ -1,5 +1,5 @@
 Wheel-Version: 1.0
-Generator: setuptools (82.0.1)
+Generator: setuptools (84.0.0)
 Root-Is-Purelib: true
 Tag: py3-none-any