stepcraft 0.1.0__py3-none-any.whl

stepcraft/__init__.py

@@ -0,0 +1,58 @@
+"""Composable function pipeline framework for Python."""
+
+from .async_runtime import (
+    HAS_RSLOOP,
+    install_rsloop,
+    rsloop_policy,
+    run_async,
+    uninstall_rsloop,
+)
+from .branching import ConditionalStep, SwitchStep
+from .constants import HAS_NUMPY, PIPE
+from .context import get_context
+from .decorators import circuit_breaker, node, piped, retry
+from .exceptions import (
+    CircuitBreakerError,
+    GraphCycleError,
+    PipelineError,
+    RetryExhaustedError,
+)
+from .fan import FanInStep, FanOutStep, MapReduceStep
+from .graph import Graph
+from .hooks import StepHook
+from .node import Node
+from .pipeline import Pipeline, PipelineBuilder
+from .pools import _POOLS, cleanup_pools, configure_pools
+from .result import ExecutionResult
+from .runtime import (
+    HAS_FREE_THREADING,
+    is_gil_enabled,
+    threads_provide_true_parallelism,
+)
+from .step import PipeStep
+from .typecheck import apply_step_beartype, beartype_enabled
+
+__all__ = [
+    # Core
+    'PIPE', 'piped', 'retry', 'circuit_breaker',
+    'PipeStep', 'Pipeline', 'PipelineBuilder',
+    # OOP
+    'Node', 'node',
+    # Graph / DAG
+    'Graph', 'ConditionalStep', 'SwitchStep',
+    # Fan-out/in
+    'FanOutStep', 'FanInStep', 'MapReduceStep',
+    # Results
+    'ExecutionResult',
+    # Observability
+    'StepHook',
+    # Shared context
+    'get_context',
+    # Errors
+    'PipelineError', 'RetryExhaustedError', 'CircuitBreakerError', 'GraphCycleError',
+    # Utilities
+    'cleanup_pools', 'configure_pools', 'HAS_NUMPY', 'HAS_RSLOOP', 'HAS_FREE_THREADING',
+    'is_gil_enabled', 'threads_provide_true_parallelism',
+    'run_async', 'install_rsloop', 'uninstall_rsloop', 'rsloop_policy',
+    'apply_step_beartype', 'beartype_enabled',
+]

stepcraft/async_runtime.py

@@ -0,0 +1,50 @@
+from __future__ import annotations
+
+import asyncio
+from contextlib import contextmanager
+from typing import Any, Coroutine, Generator, TypeVar
+
+T = TypeVar("T")
+
+try:
+    import rsloop
+
+    HAS_RSLOOP = True
+except ImportError:
+    rsloop = None  # type: ignore[assignment]
+    HAS_RSLOOP = False
+
+
+def run_async(coro: Coroutine[Any, Any, T]) -> T:
+    """Run a coroutine with rsloop when installed, otherwise stdlib asyncio."""
+    if HAS_RSLOOP:
+        return rsloop.run(coro)
+    return asyncio.run(coro)
+
+
+def install_rsloop() -> None:
+    """Install rsloop as the default asyncio event loop policy."""
+    if not HAS_RSLOOP:
+        raise ImportError(
+            "rsloop is not installed. Install with: pip install stepcraft[rsloop]"
+        )
+    rsloop.install()
+
+
+def uninstall_rsloop() -> None:
+    """Restore the previous asyncio event loop policy."""
+    if not HAS_RSLOOP:
+        raise ImportError(
+            "rsloop is not installed. Install with: pip install stepcraft[rsloop]"
+        )
+    rsloop.uninstall()
+
+
+@contextmanager
+def rsloop_policy() -> Generator[None, None, None]:
+    """Context manager that temporarily installs rsloop as the event loop policy."""
+    install_rsloop()
+    try:
+        yield
+    finally:
+        uninstall_rsloop()

stepcraft/branching.py

@@ -0,0 +1,61 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Any, Callable, Dict, Generic
+
+from .constants import R, T
+from .utils import _async_run_branch_value, _get_func_name, _run_branch_value
+
+
+@dataclass
+class ConditionalStep(Generic[T, R]):
+    """Route data through different branches based on a condition."""
+
+    condition: Callable[[T], bool]
+    if_true: Any
+    if_false: Any = None
+
+    def run(self, value: T) -> R:
+        if self.condition(value):
+            return _run_branch_value(self.if_true, value)
+        return _run_branch_value(self.if_false, value)
+
+    async def async_run(self, value: T) -> R:
+        if self.condition(value):
+            return await _async_run_branch_value(self.if_true, value)
+        return await _async_run_branch_value(self.if_false, value)
+
+    @property
+    def _func_name(self) -> str:
+        return f"ConditionalStep({_get_func_name(self.condition)})"
+
+    def __or__(self, other):
+        from .pipeline import Pipeline
+        return Pipeline([self]) | other
+
+
+@dataclass
+class SwitchStep(Generic[T, R]):
+    """Route data through one of many branches based on a key function."""
+
+    key: Callable[[T], str]
+    branches: Dict[str, Any]
+    default: Any = None
+
+    def _get_branch(self, value: T) -> Any:
+        k = self.key(value)
+        return self.branches.get(k, self.default)
+
+    def run(self, value: T) -> R:
+        return _run_branch_value(self._get_branch(value), value)
+
+    async def async_run(self, value: T) -> R:
+        return await _async_run_branch_value(self._get_branch(value), value)
+
+    @property
+    def _func_name(self) -> str:
+        return f"SwitchStep({_get_func_name(self.key)})"
+
+    def __or__(self, other):
+        from .pipeline import Pipeline
+        return Pipeline([self]) | other

stepcraft/config.py

@@ -0,0 +1,26 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from enum import Enum
+from typing import Tuple
+
+
+@dataclass(frozen=True)
+class RetryConfig:
+    attempts: int = 3
+    delay: float = 1.0
+    backoff: float = 2.0
+    errors: Tuple[type, ...] = (Exception,)
+
+
+@dataclass(frozen=True)
+class CircuitBreakerConfig:
+    threshold: int = 5
+    timeout: float = 60.0
+    half_open_max_calls: int = 1
+
+
+class CircuitState(Enum):
+    CLOSED = 0
+    OPEN = 1
+    HALF_OPEN = 2

stepcraft/constants.py

@@ -0,0 +1,16 @@
+from __future__ import annotations
+
+import logging
+from typing import TypeVar
+
+try:
+    import numpy as np
+    HAS_NUMPY = True
+except ImportError:
+    np = None  # type: ignore[assignment]
+    HAS_NUMPY = False
+
+PIPE: object = object()
+T = TypeVar("T")
+R = TypeVar("R")
+logger = logging.getLogger(__name__)

stepcraft/context.py

@@ -0,0 +1,66 @@
+from __future__ import annotations
+
+import contextvars
+import functools
+from contextlib import contextmanager
+from typing import Any, Callable, Dict, Iterator, Optional
+
+# Shared, read-only-ish context for steps to access during a pipeline run.
+_PIPELINE_CONTEXT: contextvars.ContextVar[Optional[Dict[str, Any]]] = (
+    contextvars.ContextVar("stepcraft_context", default=None)
+)
+
+
+def get_context() -> Dict[str, Any]:
+    """Return the active pipeline context, or an empty dict outside a run."""
+    ctx = _PIPELINE_CONTEXT.get()
+    return ctx if ctx is not None else {}
+
+
+def _set_context(ctx: Optional[Dict[str, Any]]) -> contextvars.Token:
+    return _PIPELINE_CONTEXT.set(ctx)
+
+
+def _reset_context(token: contextvars.Token) -> None:
+    _PIPELINE_CONTEXT.reset(token)
+
+
+@contextmanager
+def activate_context(ctx: Optional[Dict[str, Any]]) -> Iterator[None]:
+    """Make *ctx* visible to steps via ``get_context()`` for the duration."""
+    if ctx is None:
+        yield
+        return
+    token = _set_context(ctx)
+    try:
+        yield
+    finally:
+        _reset_context(token)
+
+
+def _run_with_captured_context(
+    ctx: Optional[Dict[str, Any]],
+    fn: Callable[..., Any],
+    /,
+    *args: Any,
+    **kwargs: Any,
+) -> Any:
+    """Run *fn* after setting pipeline context (for process-pool workers)."""
+    if ctx is None:
+        return fn(*args, **kwargs)
+    token = _set_context(ctx)
+    try:
+        return fn(*args, **kwargs)
+    finally:
+        _reset_context(token)
+
+
+def wrap_worker(fn: Callable[..., Any]) -> Callable[..., Any]:
+    """Wrap *fn* so pool workers inherit the caller's pipeline context."""
+    ctx = _PIPELINE_CONTEXT.get()
+    return functools.partial(_run_with_captured_context, ctx, fn)
+
+
+# Backward-compatible aliases for thread vs process dispatch sites.
+wrap_thread_worker = wrap_worker
+wrap_process_worker = wrap_worker

stepcraft/decorators.py

@@ -0,0 +1,160 @@
+from __future__ import annotations
+
+from typing import Callable, Optional, Tuple, Union
+
+from .config import CircuitBreakerConfig, RetryConfig
+from .constants import HAS_NUMPY, logger, np
+from .node import Node
+from .step import PipeStep
+from .typecheck import apply_step_beartype, resolve_output_schema
+from .utils import _get_func_name
+
+
+def piped(
+    func: Optional[Callable] = None,
+    *,
+    batch_size: int = 1,
+    parallel: Optional[str] = None,
+    auto_map: bool = True,
+    map: Optional[bool] = None,
+    timeout: Optional[float] = None,
+    cancel_on_timeout: bool = False,
+    schema: Optional[type] = None,
+    jit: bool = False,
+    vectorize: bool = False,
+) -> Union[PipeStep, Callable[[Callable], PipeStep]]:
+    """Create a PipeStep from a function.
+
+    ``map`` is a README-friendly alias for ``auto_map``; when given it wins.
+    """
+    if map is not None:
+        auto_map = map
+
+    def decorator(f: Callable) -> PipeStep:
+        if parallel and batch_size > 1:
+            logger.warning(
+                "%s: parallel=%r takes precedence over batch_size=%d; "
+                "batching is ignored when both are set",
+                _get_func_name(f), parallel, batch_size,
+            )
+
+        checked = apply_step_beartype(f)
+        optimized = checked
+        output_schema = resolve_output_schema(f, schema)
+
+        if jit:
+            try:
+                import numba
+                optimized = numba.njit(fastmath=True, cache=True, nogil=True)(optimized)
+            except ImportError:
+                logger.warning(
+                    "jit=True on %s but numba is not installed; running without JIT. "
+                    "Install with: pip install stepcraft[numba]",
+                    _get_func_name(f),
+                )
+
+        if vectorize:
+            vectorized = False
+            try:
+                import numba
+                optimized = numba.vectorize(
+                    ['float64(float64)', 'float32(float32)', 'int64(int64)'],
+                    nopython=True, cache=True,
+                )(optimized)
+                vectorized = True
+            except ImportError:
+                if HAS_NUMPY:
+                    optimized = np.vectorize(optimized, cache=True)
+                    vectorized = True
+            if not vectorized:
+                logger.warning(
+                    "vectorize=True on %s but neither numba nor numpy is available; "
+                    "running without vectorization. "
+                    "Install with: pip install stepcraft[numba] or stepcraft[numpy]",
+                    _get_func_name(f),
+                )
+
+        return PipeStep(
+            func=optimized,
+            batch_size=batch_size,
+            parallel=parallel,
+            auto_map=auto_map,
+            timeout=timeout,
+            cancel_on_timeout=cancel_on_timeout,
+            schema=output_schema,
+        )
+
+    return decorator(func) if func else decorator
+
+
+def retry(
+    *,
+    max_attempts: int = 3,
+    delay: float = 1.0,
+    backoff: float = 2.0,
+    errors: Tuple[type, ...] = (Exception,),
+) -> Callable[[PipeStep], PipeStep]:
+    """Add retry capability to a pipeline step."""
+
+    def decorator(step: PipeStep) -> PipeStep:
+        return step.copy(
+            retry_config=RetryConfig(
+                attempts=max_attempts, delay=delay, backoff=backoff, errors=errors,
+            ),
+        )
+
+    return decorator
+
+
+def circuit_breaker(
+    *,
+    threshold: int = 5,
+    timeout: float = 60.0,
+    failure_threshold: int = None,
+    recovery_timeout: float = None,
+) -> Callable[[PipeStep], PipeStep]:
+    """Add circuit breaker capability to a pipeline step."""
+    if failure_threshold is not None:
+        threshold = failure_threshold
+    if recovery_timeout is not None:
+        timeout = recovery_timeout
+
+    def decorator(step: PipeStep) -> PipeStep:
+        return step.copy(
+            circuit_config=CircuitBreakerConfig(threshold=threshold, timeout=timeout),
+        )
+
+    return decorator
+
+
+def node(
+    func: Optional[Callable] = None,
+    *,
+    setup: Optional[Callable] = None,
+    teardown: Optional[Callable] = None,
+) -> Union[Node, Callable[[Callable], Node]]:
+    """Decorator to create a Node from a function."""
+
+    def decorator(f: Callable) -> Node:
+        display_name = _get_func_name(f)
+        checked = apply_step_beartype(f)
+
+        class FuncNode(Node):
+            @property
+            def _func_name(self) -> str:
+                return display_name
+
+            def __repr__(self) -> str:
+                return f"{display_name}()"
+
+            def process(self, *args, **kwargs):
+                return checked(*args, **kwargs)
+
+        if setup:
+            FuncNode.setup = setup
+        if teardown:
+            FuncNode.teardown = teardown
+
+        return FuncNode()
+
+    return decorator(func) if func else decorator

stepcraft/exceptions.py

@@ -0,0 +1,27 @@
+from __future__ import annotations
+
+
+class PipelineError(Exception):
+    """Base pipeline exception with context."""
+
+    def __init__(self, func_name: str, original_error: Exception):
+        self.func_name = func_name
+        self.original_error = original_error
+        super().__init__(f"Pipeline step '{func_name}' failed: {original_error}")
+
+
+class RetryExhaustedError(PipelineError):
+    """Raised when retry attempts are exhausted."""
+
+    def __str__(self):
+        return f"RetryExhausted: {super().__str__()}"
+
+
+class CircuitBreakerError(PipelineError):
+    """Raised when circuit breaker is open."""
+    pass
+
+
+class GraphCycleError(Exception):
+    """Raised when a cycle is detected in the DAG."""
+    pass

stepcraft/fan.py

@@ -0,0 +1,153 @@
+from __future__ import annotations
+
+import asyncio
+import inspect
+import itertools
+from dataclasses import dataclass, field
+from typing import Any, Callable, Generic, Iterable, List, Optional, Tuple
+
+from .constants import R, T
+from .context import wrap_worker
+from .pools import _get_pool
+from .runtime import resolve_parallel_kind
+
+
+def _run_branch_on(branch: Any, value: Any) -> Any:
+    """Module-level (picklable) helper for ProcessPoolExecutor fan-out."""
+    return branch.run(value)
+
+
+@dataclass
+class FanOutStep(Generic[T, R]):
+    """Execute multiple branches in parallel."""
+
+    branches: Tuple[Any, ...]
+    parallel: Optional[str] = None
+
+    def __post_init__(self):
+        parallel = resolve_parallel_kind(self.parallel)
+        if parallel != self.parallel:
+            object.__setattr__(self, 'parallel', parallel)
+
+    def run(self, value: T) -> Tuple[R, ...]:
+        if self.parallel:
+            pool = _get_pool(self.parallel)
+            runner = wrap_worker(_run_branch_on)
+            return tuple(
+                pool.map(runner, self.branches, itertools.repeat(value))
+            )
+        return tuple(branch.run(value) for branch in self.branches)
+
+    async def async_run(self, value: T) -> Tuple[R, ...]:
+        if self.parallel:
+            pool = _get_pool(self.parallel)
+            loop = asyncio.get_running_loop()
+            runner = wrap_worker(_run_branch_on)
+            tasks = [
+                loop.run_in_executor(pool, runner, branch, value)
+                for branch in self.branches
+            ]
+            return tuple(await asyncio.gather(*tasks))
+
+        tasks = []
+        for branch in self.branches:
+            if hasattr(branch, 'async_run'):
+                tasks.append(branch.async_run(value))
+            else:
+                loop = asyncio.get_running_loop()
+                call = wrap_worker(lambda b=branch, v=value: b.run(v))
+                tasks.append(loop.run_in_executor(None, call))
+        return tuple(await asyncio.gather(*tasks))
+
+    @property
+    def _func_name(self) -> str:
+        return "FanOutStep"
+
+    def __or__(self, other):
+        from .pipeline import Pipeline
+        return Pipeline([self]) | other
+
+
+@dataclass
+class FanInStep(Generic[T, R]):
+    """Combine multiple inputs into single output."""
+
+    combiner: Callable[..., R]
+
+    def run(self, values: Tuple[T, ...]) -> R:
+        return self.combiner(*values)
+
+    async def async_run(self, values: Tuple[T, ...]) -> R:
+        result = self.combiner(*values)
+        if asyncio.iscoroutine(result):
+            return await result
+        return result
+
+    @property
+    def _func_name(self) -> str:
+        return "FanInStep"
+
+    def __or__(self, other):
+        from .pipeline import Pipeline
+        return Pipeline([self]) | other
+
+
+@dataclass
+class MapReduceStep(Generic[T, R]):
+    """Map-reduce with optional batching."""
+
+    mapper: Callable[[T], Any]
+    reducer: Callable[[Iterable[Any]], R]
+    batch_size: int = 1
+    _mapper_is_async: bool = field(default=False, init=False)
+
+    def __post_init__(self):
+        object.__setattr__(
+            self, '_mapper_is_async', inspect.iscoroutinefunction(self.mapper),
+        )
+
+    async def _map_item(self, item: T) -> Any:
+        if self._mapper_is_async:
+            return await self.mapper(item)
+        result = self.mapper(item)
+        if asyncio.iscoroutine(result):
+            return await result
+        return result
+
+    async def _map_batch(self, batch: List[T]) -> List[Any]:
+        return list(await asyncio.gather(*(self._map_item(x) for x in batch)))
+
+    def run(self, items: Iterable[T]) -> R:
+        results: List[Any] = []
+        batch: List[T] = []
+        for item in items:
+            batch.append(item)
+            if len(batch) == self.batch_size:
+                results.extend(map(self.mapper, batch))
+                batch.clear()
+        if batch:
+            results.extend(map(self.mapper, batch))
+        return self.reducer(results)
+
+    async def async_run(self, items: Iterable[T]) -> R:
+        results: List[Any] = []
+        batch: List[T] = []
+        for item in items:
+            batch.append(item)
+            if len(batch) == self.batch_size:
+                results.extend(await self._map_batch(batch))
+                batch.clear()
+        if batch:
+            results.extend(await self._map_batch(batch))
+        out = self.reducer(results)
+        if asyncio.iscoroutine(out):
+            return await out
+        return out
+
+    @property
+    def _func_name(self) -> str:
+        return "MapReduceStep"
+
+    def __or__(self, other):
+        from .pipeline import Pipeline
+        return Pipeline([self]) | other