feat: registration for standalone functions

Author: mrkaye97

examples/simple/worker.py

@@ -8,12 +8,13 @@
 hatchet = Hatchet(debug=True)
 
 
-@hatchet.function(timeout="11s")
+@hatchet.function()
 def step1(context: Context) -> dict[str, str]:
-    print("executed step1")
-    return {
-        "step1": "step1",
-    }
+    message = "Hello from Hatchet!"
+
+    context.log(message)
+
+    return {"message": message}
 
 
 def main() -> None:

hatchet_sdk/v2/hatchet.py

@@ -4,10 +4,12 @@
     TYPE_CHECKING,
     Any,
     Callable,
+    Generic,
     Optional,
     ParamSpec,
     Type,
     TypeVar,
+    Union,
     cast,
 )
 
@@ -55,6 +57,56 @@ def transform_desired_worker_label(d: DesiredWorkerLabel) -> DesiredWorkerLabels
     )
 
 
+class Function(Generic[R, TWorkflowInput]):
+    def __init__(
+        self,
+        fn: Callable[[Context], R],
+        hatchet: "Hatchet",
+        name: str = "",
+        on_events: list[str] = [],
+        on_crons: list[str] = [],
+        version: str = "",
+        timeout: str = "60m",
+        schedule_timeout: str = "5m",
+        sticky: StickyStrategy | None = None,
+        retries: int = 0,
+        rate_limits: list[RateLimit] = [],
+        desired_worker_labels: dict[str, DesiredWorkerLabel] = {},
+        concurrency: ConcurrencyExpression | None = None,
+        on_failure: Union["Function[R]", None] = None,
+        default_priority: int = 1,
+        input_validator: Type[TWorkflowInput] | None = None,
+        backoff_factor: float | None = None,
+        backoff_max_seconds: int | None = None,
+    ) -> None:
+        def func(_: Any, context: Context) -> R:
+            return fn(context)
+
+        self.hatchet = hatchet
+        self.step: Step[R] = hatchet.step(
+            name=name or fn.__name__,
+            timeout=timeout,
+            retries=retries,
+            rate_limits=rate_limits,
+            desired_worker_labels=desired_worker_labels,
+            backoff_factor=backoff_factor,
+            backoff_max_seconds=backoff_max_seconds,
+        )(func)
+        self.on_failure_step = on_failure
+        self.workflow_config = WorkflowConfig(
+            name=name or fn.__name__,
+            on_events=on_events,
+            on_crons=on_crons,
+            version=version,
+            timeout=timeout,
+            schedule_timeout=schedule_timeout,
+            sticky=sticky,
+            default_priority=default_priority,
+            concurrency=concurrency,
+            input_validator=input_validator or cast(Type[TWorkflowInput], EmptyModel),
+        )
+
+
 class Hatchet:
     """
     Main client for interacting with the Hatchet SDK.
@@ -207,44 +259,38 @@ def function(
         timeout: str = "60m",
         schedule_timeout: str = "5m",
         sticky: StickyStrategy | None = None,
-        default_priority: int = 1,
+        retries: int = 0,
+        rate_limits: list[RateLimit] = [],
+        desired_worker_labels: dict[str, DesiredWorkerLabel] = {},
         concurrency: ConcurrencyExpression | None = None,
+        on_failure: Union["Function[Any]", None] = None,
+        default_priority: int = 1,
         input_validator: Type[TWorkflowInput] | None = None,
-    ) -> Callable[[Callable[[Context], R]], BaseWorkflowImpl]:
-        def inner(func: Callable[[Context], R]) -> BaseWorkflowImpl:
-            declaration = WorkflowDeclaration[TWorkflowInput](
-                WorkflowConfig(
-                    name=name or func.__name__,
-                    on_events=on_events,
-                    on_crons=on_crons,
-                    version=version,
-                    timeout=timeout,
-                    schedule_timeout=schedule_timeout,
-                    sticky=sticky,
-                    default_priority=default_priority,
-                    concurrency=concurrency,
-                    input_validator=input_validator
-                    or cast(Type[TWorkflowInput], EmptyModel),
-                ),
-                self,
+        backoff_factor: float | None = None,
+        backoff_max_seconds: int | None = None,
+    ) -> Callable[[Callable[[Context], R]], Function[R, TWorkflowInput]]:
+        def inner(func: Callable[[Context], R]) -> Function[R, TWorkflowInput]:
+            return Function[R, TWorkflowInput](
+                func,
+                hatchet=self,
+                name=name,
+                on_events=on_events,
+                on_crons=on_crons,
+                version=version,
+                timeout=timeout,
+                schedule_timeout=schedule_timeout,
+                sticky=sticky,
+                retries=retries,
+                rate_limits=rate_limits,
+                desired_worker_labels=desired_worker_labels,
+                concurrency=concurrency,
+                on_failure=on_failure,
+                default_priority=default_priority,
+                input_validator=input_validator,
+                backoff_factor=backoff_factor,
+                backoff_max_seconds=backoff_max_seconds,
             )
 
-            class Workflow(BaseWorkflowImpl):
-                config = declaration.config
-
-                @self.step(
-                    name=declaration.config.name,
-                    timeout=timeout,
-                    retries=0,
-                    rate_limits=[],
-                    backoff_factor=None,
-                    backoff_max_seconds=None,
-                )
-                def fn(self, context: Context) -> R:
-                    return func(context)
-
-            return Workflow()
-
         return inner
 
     def worker(

hatchet_sdk/worker/worker.py

@@ -37,6 +37,7 @@
 
 if TYPE_CHECKING:
     from hatchet_sdk.v2 import BaseWorkflowImpl
+    from hatchet_sdk.v2.hatchet import Function
 
 T = TypeVar("T")
 
@@ -110,8 +111,31 @@ def register_workflow_from_opts(
             logger.error(e)
             sys.exit(1)
 
-    def register_function(self, function: "BaseWorkflowImpl") -> None:
-        self.register_workflow(function)
+    def register_function(self, function: "Function[Any]") -> None:
+        from hatchet_sdk.v2.workflows import BaseWorkflowImpl
+
+        declaration = function.hatchet.declare_workflow(
+            **function.workflow_config.model_dump()
+        )
+
+        class Workflow(BaseWorkflowImpl):
+            config = declaration.config
+
+            @property
+            def default_steps(self) -> list[Step[Any]]:
+                return [function.step]
+
+            @property
+            def on_failure_steps(self) -> list[Step[Any]]:
+                if not function.on_failure_step:
+                    return []
+
+                step = function.on_failure_step.step
+                step.type = StepType.ON_FAILURE
+
+                return [step]
+
+        self.register_workflow(Workflow())
 
     def register_workflow(self, workflow: Union["BaseWorkflowImpl", Any]) -> None:
         namespace = self.client.config.namespace
@@ -125,7 +149,9 @@ def register_workflow(self, workflow: Union["BaseWorkflowImpl", Any]) -> None:
             logger.error(e)
             sys.exit(1)
 
+        print(workflow.steps)
         for step in workflow.steps:
+            print(step)
             action_name = workflow.create_action_name(namespace, step)
             self.action_registry[action_name] = step
             return_type = get_type_hints(step.fn).get("return")