hatchet-sdk 0.42.2__py3-none-any.whl → 0.42.4__py3-none-any.whl

hatchet_sdk/clients/rest/models/events.py

@@ -0,0 +1,110 @@
+# coding: utf-8
+
+"""
+    Hatchet API
+
+    The Hatchet API
+
+    The version of the OpenAPI document: 1.0.0
+    Generated by OpenAPI Generator (https://openapi-generator.tech)
+
+    Do not edit the class manually.
+"""  # noqa: E501
+
+
+from __future__ import annotations
+
+import json
+import pprint
+import re  # noqa: F401
+from typing import Any, ClassVar, Dict, List, Optional, Set
+
+from pydantic import BaseModel, ConfigDict, Field
+from typing_extensions import Self
+
+from hatchet_sdk.clients.rest.models.api_resource_meta import APIResourceMeta
+from hatchet_sdk.clients.rest.models.event import Event
+
+
+class Events(BaseModel):
+    """
+    Events
+    """  # noqa: E501
+
+    metadata: APIResourceMeta
+    events: List[Event] = Field(description="The events.")
+    __properties: ClassVar[List[str]] = ["metadata", "events"]
+
+    model_config = ConfigDict(
+        populate_by_name=True,
+        validate_assignment=True,
+        protected_namespaces=(),
+    )
+
+    def to_str(self) -> str:
+        """Returns the string representation of the model using alias"""
+        return pprint.pformat(self.model_dump(by_alias=True))
+
+    def to_json(self) -> str:
+        """Returns the JSON representation of the model using alias"""
+        # TODO: pydantic v2: use .model_dump_json(by_alias=True, exclude_unset=True) instead
+        return json.dumps(self.to_dict())
+
+    @classmethod
+    def from_json(cls, json_str: str) -> Optional[Self]:
+        """Create an instance of Events from a JSON string"""
+        return cls.from_dict(json.loads(json_str))
+
+    def to_dict(self) -> Dict[str, Any]:
+        """Return the dictionary representation of the model using alias.
+
+        This has the following differences from calling pydantic's
+        `self.model_dump(by_alias=True)`:
+
+        * `None` is only added to the output dict for nullable fields that
+          were set at model initialization. Other fields with value `None`
+          are ignored.
+        """
+        excluded_fields: Set[str] = set([])
+
+        _dict = self.model_dump(
+            by_alias=True,
+            exclude=excluded_fields,
+            exclude_none=True,
+        )
+        # override the default output from pydantic by calling `to_dict()` of metadata
+        if self.metadata:
+            _dict["metadata"] = self.metadata.to_dict()
+        # override the default output from pydantic by calling `to_dict()` of each item in events (list)
+        _items = []
+        if self.events:
+            for _item_events in self.events:
+                if _item_events:
+                    _items.append(_item_events.to_dict())
+            _dict["events"] = _items
+        return _dict
+
+    @classmethod
+    def from_dict(cls, obj: Optional[Dict[str, Any]]) -> Optional[Self]:
+        """Create an instance of Events from a dict"""
+        if obj is None:
+            return None
+
+        if not isinstance(obj, dict):
+            return cls.model_validate(obj)
+
+        _obj = cls.model_validate(
+            {
+                "metadata": (
+                    APIResourceMeta.from_dict(obj["metadata"])
+                    if obj.get("metadata") is not None
+                    else None
+                ),
+                "events": (
+                    [Event.from_dict(_item) for _item in obj["events"]]
+                    if obj.get("events") is not None
+                    else None
+                ),
+            }
+        )
+        return _obj

hatchet_sdk/clients/rest/models/scheduled_workflows.py

@@ -20,10 +20,13 @@ import re  # noqa: F401
 from datetime import datetime
 from typing import Any, ClassVar, Dict, List, Optional, Set
 
-from pydantic import BaseModel, ConfigDict, Field, StrictStr, field_validator
+from pydantic import BaseModel, ConfigDict, Field, StrictStr
 from typing_extensions import Annotated, Self
 
 from hatchet_sdk.clients.rest.models.api_resource_meta import APIResourceMeta
+from hatchet_sdk.clients.rest.models.scheduled_workflows_method import (
+    ScheduledWorkflowsMethod,
+)
 from hatchet_sdk.clients.rest.models.workflow_run_status import WorkflowRunStatus
 
 
@@ -54,7 +57,7 @@ class ScheduledWorkflows(BaseModel):
     workflow_run_id: Optional[
         Annotated[str, Field(min_length=36, strict=True, max_length=36)]
     ] = Field(default=None, alias="workflowRunId")
-    method: StrictStr
+    method: ScheduledWorkflowsMethod
     __properties: ClassVar[List[str]] = [
         "metadata",
         "tenantId",
@@ -71,13 +74,6 @@ class ScheduledWorkflows(BaseModel):
         "method",
     ]
 
-    @field_validator("method")
-    def method_validate_enum(cls, value):
-        """Validates the enum"""
-        if value not in set(["DEFAULT", "API"]):
-            raise ValueError("must be one of enum values ('DEFAULT', 'API')")
-        return value
-
     model_config = ConfigDict(
         populate_by_name=True,
         validate_assignment=True,

hatchet_sdk/clients/rest/models/scheduled_workflows_method.py

@@ -0,0 +1,37 @@
+# coding: utf-8
+
+"""
+    Hatchet API
+
+    The Hatchet API
+
+    The version of the OpenAPI document: 1.0.0
+    Generated by OpenAPI Generator (https://openapi-generator.tech)
+
+    Do not edit the class manually.
+"""  # noqa: E501
+
+
+from __future__ import annotations
+
+import json
+from enum import Enum
+
+from typing_extensions import Self
+
+
+class ScheduledWorkflowsMethod(str, Enum):
+    """
+    ScheduledWorkflowsMethod
+    """
+
+    """
+    allowed enum values
+    """
+    DEFAULT = "DEFAULT"
+    API = "API"
+
+    @classmethod
+    def from_json(cls, json_str: str) -> Self:
+        """Create an instance of ScheduledWorkflowsMethod from a JSON string"""
+        return cls(json.loads(json_str))

hatchet_sdk/clients/rest/models/worker.py

@@ -28,6 +28,7 @@ from hatchet_sdk.clients.rest.models.recent_step_runs import RecentStepRuns
 from hatchet_sdk.clients.rest.models.semaphore_slots import SemaphoreSlots
 from hatchet_sdk.clients.rest.models.worker_label import WorkerLabel
 from hatchet_sdk.clients.rest.models.worker_runtime_info import WorkerRuntimeInfo
+from hatchet_sdk.clients.rest.models.worker_type import WorkerType
 
 
 class Worker(BaseModel):
@@ -37,7 +38,7 @@ class Worker(BaseModel):
 
     metadata: APIResourceMeta
     name: StrictStr = Field(description="The name of the worker.")
-    type: StrictStr
+    type: WorkerType
     last_heartbeat_at: Optional[datetime] = Field(
         default=None,
         description="The time this worker last sent a heartbeat.",
@@ -108,15 +109,6 @@ class Worker(BaseModel):
         "runtimeInfo",
     ]
 
-    @field_validator("type")
-    def type_validate_enum(cls, value):
-        """Validates the enum"""
-        if value not in set(["SELFHOSTED", "MANAGED", "WEBHOOK"]):
-            raise ValueError(
-                "must be one of enum values ('SELFHOSTED', 'MANAGED', 'WEBHOOK')"
-            )
-        return value
-
     @field_validator("status")
     def status_validate_enum(cls, value):
         """Validates the enum"""

hatchet_sdk/clients/rest/models/worker_type.py

@@ -0,0 +1,38 @@
+# coding: utf-8
+
+"""
+    Hatchet API
+
+    The Hatchet API
+
+    The version of the OpenAPI document: 1.0.0
+    Generated by OpenAPI Generator (https://openapi-generator.tech)
+
+    Do not edit the class manually.
+"""  # noqa: E501
+
+
+from __future__ import annotations
+
+import json
+from enum import Enum
+
+from typing_extensions import Self
+
+
+class WorkerType(str, Enum):
+    """
+    WorkerType
+    """
+
+    """
+    allowed enum values
+    """
+    SELFHOSTED = "SELFHOSTED"
+    MANAGED = "MANAGED"
+    WEBHOOK = "WEBHOOK"
+
+    @classmethod
+    def from_json(cls, json_str: str) -> Self:
+        """Create an instance of WorkerType from a JSON string"""
+        return cls(json.loads(json_str))

hatchet_sdk/clients/rest/models/workflow_concurrency.py

@@ -19,9 +19,13 @@ import pprint
 import re  # noqa: F401
 from typing import Any, ClassVar, Dict, List, Optional, Set
 
-from pydantic import BaseModel, ConfigDict, Field, StrictInt, StrictStr, field_validator
+from pydantic import BaseModel, ConfigDict, Field, StrictInt, StrictStr
 from typing_extensions import Self
 
+from hatchet_sdk.clients.rest.models.concurrency_limit_strategy import (
+    ConcurrencyLimitStrategy,
+)
+
 
 class WorkflowConcurrency(BaseModel):
     """
@@ -31,7 +35,7 @@ class WorkflowConcurrency(BaseModel):
     max_runs: StrictInt = Field(
         description="The maximum number of concurrent workflow runs.", alias="maxRuns"
     )
-    limit_strategy: StrictStr = Field(
+    limit_strategy: ConcurrencyLimitStrategy = Field(
         description="The strategy to use when the concurrency limit is reached.",
         alias="limitStrategy",
     )
@@ -45,17 +49,6 @@ class WorkflowConcurrency(BaseModel):
         "getConcurrencyGroup",
     ]
 
-    @field_validator("limit_strategy")
-    def limit_strategy_validate_enum(cls, value):
-        """Validates the enum"""
-        if value not in set(
-            ["CANCEL_IN_PROGRESS", "DROP_NEWEST", "QUEUE_NEWEST", "GROUP_ROUND_ROBIN"]
-        ):
-            raise ValueError(
-                "must be one of enum values ('CANCEL_IN_PROGRESS', 'DROP_NEWEST', 'QUEUE_NEWEST', 'GROUP_ROUND_ROBIN')"
-            )
-        return value
-
     model_config = ConfigDict(
         populate_by_name=True,
         validate_assignment=True,

hatchet_sdk/clients/rest/models/workflow_list.py

@@ -32,10 +32,10 @@ class WorkflowList(BaseModel):
     WorkflowList
     """  # noqa: E501
 
-    metadata: Optional[APIResourceMeta] = None
-    rows: Optional[List[Workflow]] = None
-    pagination: Optional[PaginationResponse] = None
-    __properties: ClassVar[List[str]] = ["metadata", "rows", "pagination"]
+    metadata: APIResourceMeta | None = None
+    rows: list[Workflow] | None = None
+    pagination: PaginationResponse | None = None
+    __properties: ClassVar[list[str]] = ["metadata", "rows", "pagination"]
 
     model_config = ConfigDict(
         populate_by_name=True,

hatchet_sdk/clients/rest/models/workflow_run.py

@@ -39,28 +39,28 @@ class WorkflowRun(BaseModel):
     metadata: APIResourceMeta
     tenant_id: StrictStr = Field(alias="tenantId")
     workflow_version_id: StrictStr = Field(alias="workflowVersionId")
-    workflow_version: Optional[WorkflowVersion] = Field(
+    workflow_version: WorkflowVersion | None = Field(
         default=None, alias="workflowVersion"
     )
     status: WorkflowRunStatus
-    display_name: Optional[StrictStr] = Field(default=None, alias="displayName")
-    job_runs: Optional[List[JobRun]] = Field(default=None, alias="jobRuns")
+    display_name: StrictStr | None = Field(default=None, alias="displayName")
+    job_runs: list[JobRun] | None = Field(default=None, alias="jobRuns")
     triggered_by: WorkflowRunTriggeredBy = Field(alias="triggeredBy")
-    input: Optional[Dict[str, Any]] = None
-    error: Optional[StrictStr] = None
-    started_at: Optional[datetime] = Field(default=None, alias="startedAt")
-    finished_at: Optional[datetime] = Field(default=None, alias="finishedAt")
-    duration: Optional[StrictInt] = None
+    input: dict[str, Any] | None = None
+    error: StrictStr | None = None
+    started_at: datetime | None = Field(default=None, alias="startedAt")
+    finished_at: datetime | None = Field(default=None, alias="finishedAt")
+    duration: StrictInt | None = None
     parent_id: Optional[
         Annotated[str, Field(min_length=36, strict=True, max_length=36)]
     ] = Field(default=None, alias="parentId")
     parent_step_run_id: Optional[
         Annotated[str, Field(min_length=36, strict=True, max_length=36)]
     ] = Field(default=None, alias="parentStepRunId")
-    additional_metadata: Optional[Dict[str, Any]] = Field(
+    additional_metadata: dict[str, Any] | None = Field(
         default=None, alias="additionalMetadata"
     )
-    __properties: ClassVar[List[str]] = [
+    __properties: ClassVar[list[str]] = [
         "metadata",
         "tenantId",
         "workflowVersionId",

hatchet_sdk/context/context.py

@@ -13,8 +13,8 @@ from hatchet_sdk.clients.rest_client import RestApi
 from hatchet_sdk.clients.run_event_listener import RunEventListenerClient
 from hatchet_sdk.clients.workflow_listener import PooledWorkflowRunListener
 from hatchet_sdk.context.worker_context import WorkerContext
-from hatchet_sdk.contracts.dispatcher_pb2 import OverridesData  # type: ignore
-from hatchet_sdk.contracts.workflows_pb2 import (  # type: ignore[attr-defined]
+from hatchet_sdk.contracts.dispatcher_pb2 import OverridesData
+from hatchet_sdk.contracts.workflows_pb2 import (
     BulkTriggerWorkflowRequest,
     TriggerWorkflowRequest,
 )
@@ -69,7 +69,7 @@ class BaseContext:
             meta = options["additional_metadata"]
 
         ## TODO: Pydantic here to simplify this
-        trigger_options: TriggerWorkflowOptions = {  # type: ignore[typeddict-item]
+        trigger_options: TriggerWorkflowOptions = {
             "parent_id": workflow_run_id,
             "parent_step_run_id": step_run_id,
             "child_key": key,
@@ -149,8 +149,7 @@ class ContextAioImpl(BaseContext):
             key = child_workflow_run.get("key")
             options = child_workflow_run.get("options", {})
 
-            ## TODO: figure out why this is failing
-            trigger_options = self._prepare_workflow_options(key, options, worker_id)  # type: ignore[arg-type]
+            trigger_options = self._prepare_workflow_options(key, options, worker_id)
 
             bulk_trigger_workflow_runs.append(
                 WorkflowRunDict(
@@ -238,14 +237,17 @@ class Context(BaseContext):
             self.input = self.data.get("input", {})
 
     def step_output(self, step: str) -> dict[str, Any] | BaseModel:
-        validators = self.validator_registry.get(step)
+        workflow_validator = next(
+            (v for k, v in self.validator_registry.items() if k.split(":")[-1] == step),
+            None,
+        )
 
         try:
             parent_step_data = cast(dict[str, Any], self.data["parents"][step])
         except KeyError:
             raise ValueError(f"Step output for '{step}' not found")
 
-        if validators and (v := validators.step_output):
+        if workflow_validator and (v := workflow_validator.step_output):
             return v.model_validate(parent_step_data)
 
         return parent_step_data

hatchet_sdk/context/worker_context.py

@@ -2,7 +2,7 @@ from hatchet_sdk.clients.dispatcher.dispatcher import DispatcherClient
 
 
 class WorkerContext:
-    _worker_id: str = None
+    _worker_id: str | None = None
     _registered_workflow_names: list[str] = []
     _labels: dict[str, str | int] = {}
 
@@ -10,18 +10,18 @@ class WorkerContext:
         self._labels = labels
         self.client = client
 
-    def labels(self):
+    def labels(self) -> dict[str, str | int]:
         return self._labels
 
-    def upsert_labels(self, labels: dict[str, str | int]):
+    def upsert_labels(self, labels: dict[str, str | int]) -> None:
         self.client.upsert_worker_labels(self._worker_id, labels)
         self._labels.update(labels)
 
-    async def async_upsert_labels(self, labels: dict[str, str | int]):
+    async def async_upsert_labels(self, labels: dict[str, str | int]) -> None:
         await self.client.async_upsert_worker_labels(self._worker_id, labels)
         self._labels.update(labels)
 
-    def id(self) -> str:
+    def id(self) -> str | None:
         return self._worker_id
 
     # def has_workflow(self, workflow_name: str):

hatchet_sdk/contracts/dispatcher_pb2.pyi

@@ -1,5 +1,3 @@
-# type: ignore
-
 from google.protobuf import timestamp_pb2 as _timestamp_pb2
 from google.protobuf.internal import containers as _containers
 from google.protobuf.internal import enum_type_wrapper as _enum_type_wrapper

hatchet_sdk/contracts/events_pb2.pyi

@@ -1,5 +1,3 @@
-# type: ignore
-
 from google.protobuf import timestamp_pb2 as _timestamp_pb2
 from google.protobuf.internal import containers as _containers
 from google.protobuf import descriptor as _descriptor

hatchet_sdk/contracts/workflows_pb2.pyi

@@ -1,5 +1,3 @@
-# type: ignore
-
 from google.protobuf import timestamp_pb2 as _timestamp_pb2
 from google.protobuf.internal import containers as _containers
 from google.protobuf.internal import enum_type_wrapper as _enum_type_wrapper

hatchet_sdk/hatchet.py

@@ -1,15 +1,13 @@
 import asyncio
 import logging
-from typing import Any, Callable, Optional, Type, TypeVar, cast, get_type_hints
+from typing import Any, Callable, Optional, Type, TypeVar, Union
 
 from pydantic import BaseModel
 from typing_extensions import deprecated
 
 from hatchet_sdk.clients.rest_client import RestApi
-
-## TODO: These type stubs need to be updated to mass MyPy, and then we can remove this ignore
-## There are file-level type ignore lines in the corresponding .pyi files.
-from hatchet_sdk.contracts.workflows_pb2 import (  # type: ignore[attr-defined]
+from hatchet_sdk.context.context import Context
+from hatchet_sdk.contracts.workflows_pb2 import (
     ConcurrencyLimitStrategy,
     CreateStepRateLimit,
     DesiredWorkerLabels,
@@ -20,6 +18,7 @@ from hatchet_sdk.features.scheduled import ScheduledClient
 from hatchet_sdk.labels import DesiredWorkerLabel
 from hatchet_sdk.loader import ClientConfig, ConfigLoader
 from hatchet_sdk.rate_limit import RateLimit
+from hatchet_sdk.v2.callable import HatchetCallable
 
 from .client import Client, new_client, new_client_raw
 from .clients.admin import AdminClient
@@ -36,6 +35,7 @@ from .workflow import (
 )
 
 T = TypeVar("T", bound=BaseModel)
+TWorkflow = TypeVar("TWorkflow", bound=object)
 
 
 def workflow(
@@ -45,30 +45,35 @@ def workflow(
     version: str = "",
     timeout: str = "60m",
     schedule_timeout: str = "5m",
-    sticky: StickyStrategy = None,
+    sticky: Union[StickyStrategy.Value, None] = None,  # type: ignore[name-defined]
     default_priority: int | None = None,
     concurrency: ConcurrencyExpression | None = None,
     input_validator: Type[T] | None = None,
-) -> Callable[[Type[WorkflowInterface]], WorkflowMeta]:
+) -> Callable[[Type[TWorkflow]], WorkflowMeta]:
     on_events = on_events or []
     on_crons = on_crons or []
 
-    def inner(cls: Type[WorkflowInterface]) -> WorkflowMeta:
-        cls.on_events = on_events
-        cls.on_crons = on_crons
-        cls.name = name or str(cls.__name__)
-        cls.version = version
-        cls.timeout = timeout
-        cls.schedule_timeout = schedule_timeout
-        cls.sticky = sticky
-        cls.default_priority = default_priority
-        cls.concurrency_expression = concurrency
+    def inner(cls: Type[TWorkflow]) -> WorkflowMeta:
+        nonlocal name
+        name = name or str(cls.__name__)
+
+        setattr(cls, "on_events", on_events)
+        setattr(cls, "on_crons", on_crons)
+        setattr(cls, "name", name)
+        setattr(cls, "version", version)
+        setattr(cls, "timeout", timeout)
+        setattr(cls, "schedule_timeout", schedule_timeout)
+        setattr(cls, "sticky", sticky)
+        setattr(cls, "default_priority", default_priority)
+        setattr(cls, "concurrency_expression", concurrency)
+
         # Define a new class with the same name and bases as the original, but
         # with WorkflowMeta as its metaclass
 
         ## TODO: Figure out how to type this metaclass correctly
-        cls.input_validator = input_validator
-        return WorkflowMeta(cls.name, cls.__bases__, dict(cls.__dict__))
+        setattr(cls, "input_validator", input_validator)
+
+        return WorkflowMeta(name, cls.__bases__, dict(cls.__dict__))
 
     return inner
 
@@ -82,37 +87,39 @@ def step(
     desired_worker_labels: dict[str, DesiredWorkerLabel] = {},
     backoff_factor: float | None = None,
     backoff_max_seconds: int | None = None,
-) -> Callable[[WorkflowStepProtocol], WorkflowStepProtocol]:
+) -> Callable[..., Any]:
     parents = parents or []
 
-    def inner(func: WorkflowStepProtocol) -> WorkflowStepProtocol:
+    def inner(func: Callable[[Context], Any]) -> Callable[[Context], Any]:
         limits = None
         if rate_limits:
-            limits = [
-                CreateStepRateLimit(key=rate_limit.key, units=rate_limit.units)
-                for rate_limit in rate_limits or []
-            ]
-
-        func._step_name = name.lower() or str(func.__name__).lower()
-        func._step_parents = parents
-        func._step_timeout = timeout
-        func._step_retries = retries
-        func._step_rate_limits = limits
-        func._step_backoff_factor = backoff_factor
-        func._step_backoff_max_seconds = backoff_max_seconds
-
-        func._step_desired_worker_labels = {}
-
-        for key, d in desired_worker_labels.items():
+            limits = [rate_limit._req for rate_limit in rate_limits or []]
+
+        setattr(func, "_step_name", name.lower() or str(func.__name__).lower())
+        setattr(func, "_step_parents", parents)
+        setattr(func, "_step_timeout", timeout)
+        setattr(func, "_step_retries", retries)
+        setattr(func, "_step_rate_limits", retries)
+        setattr(func, "_step_rate_limits", limits)
+        setattr(func, "_step_backoff_factor", backoff_factor)
+        setattr(func, "_step_backoff_max_seconds", backoff_max_seconds)
+
+        def create_label(d: DesiredWorkerLabel) -> DesiredWorkerLabels:
             value = d["value"] if "value" in d else None
-            func._step_desired_worker_labels[key] = DesiredWorkerLabels(
+            return DesiredWorkerLabels(
                 strValue=str(value) if not isinstance(value, int) else None,
                 intValue=value if isinstance(value, int) else None,
-                required=d["required"] if "required" in d else None,
+                required=d["required"] if "required" in d else None,  # type: ignore[arg-type]
                 weight=d["weight"] if "weight" in d else None,
-                comparator=d["comparator"] if "comparator" in d else None,
+                comparator=d["comparator"] if "comparator" in d else None,  # type: ignore[arg-type]
             )
 
+        setattr(
+            func,
+            "_step_desired_worker_labels",
+            {key: create_label(d) for key, d in desired_worker_labels.items()},
+        )
+
         return func
 
     return inner
@@ -125,21 +132,23 @@ def on_failure_step(
     rate_limits: list[RateLimit] | None = None,
     backoff_factor: float | None = None,
     backoff_max_seconds: int | None = None,
-) -> Callable[[WorkflowStepProtocol], WorkflowStepProtocol]:
-    def inner(func: WorkflowStepProtocol) -> WorkflowStepProtocol:
+) -> Callable[..., Any]:
+    def inner(func: Callable[[Context], Any]) -> Callable[[Context], Any]:
         limits = None
         if rate_limits:
             limits = [
-                CreateStepRateLimit(key=rate_limit.key, units=rate_limit.units)
+                CreateStepRateLimit(key=rate_limit.static_key, units=rate_limit.units)  # type: ignore[arg-type]
                 for rate_limit in rate_limits or []
             ]
 
-        func._on_failure_step_name = name.lower() or str(func.__name__).lower()
-        func._on_failure_step_timeout = timeout
-        func._on_failure_step_retries = retries
-        func._on_failure_step_rate_limits = limits
-        func._on_failure_step_backoff_factor = backoff_factor
-        func._on_failure_step_backoff_max_seconds = backoff_max_seconds
+        setattr(
+            func, "_on_failure_step_name", name.lower() or str(func.__name__).lower()
+        )
+        setattr(func, "_on_failure_step_timeout", timeout)
+        setattr(func, "_on_failure_step_retries", retries)
+        setattr(func, "_on_failure_step_rate_limits", limits)
+        setattr(func, "_on_failure_step_backoff_factor", backoff_factor)
+        setattr(func, "_on_failure_step_backoff_max_seconds", backoff_max_seconds)
 
         return func
 
@@ -150,11 +159,15 @@ def concurrency(
     name: str = "",
     max_runs: int = 1,
     limit_strategy: ConcurrencyLimitStrategy = ConcurrencyLimitStrategy.CANCEL_IN_PROGRESS,
-) -> Callable[[WorkflowStepProtocol], WorkflowStepProtocol]:
-    def inner(func: WorkflowStepProtocol) -> WorkflowStepProtocol:
-        func._concurrency_fn_name = name.lower() or str(func.__name__).lower()
-        func._concurrency_max_runs = max_runs
-        func._concurrency_limit_strategy = limit_strategy
+) -> Callable[..., Any]:
+    def inner(func: Callable[[Context], Any]) -> Callable[[Context], Any]:
+        setattr(
+            func,
+            "_concurrency_fn_name",
+            name.lower() or str(func.__name__).lower(),
+        )
+        setattr(func, "_concurrency_max_runs", max_runs)
+        setattr(func, "_concurrency_limit_strategy", limit_strategy)
 
         return func
 

hatchet_sdk/labels.py

@@ -1,7 +1,7 @@
 from typing import TypedDict
 
 
-class DesiredWorkerLabel(TypedDict):
+class DesiredWorkerLabel(TypedDict, total=False):
     value: str | int
     required: bool | None = None
     weight: int | None = None