diff --git a/task_sdk/src/airflow/sdk/api/client.py b/task_sdk/src/airflow/sdk/api/client.py index c740dceb01041..f52e6dde6227b 100644 --- a/task_sdk/src/airflow/sdk/api/client.py +++ b/task_sdk/src/airflow/sdk/api/client.py @@ -19,10 +19,9 @@ import sys import uuid -from typing import TYPE_CHECKING, Any +from typing import TYPE_CHECKING, Any, TypeVar import httpx -import methodtools import msgspec import structlog from pydantic import BaseModel @@ -42,6 +41,19 @@ if TYPE_CHECKING: from datetime import datetime + from airflow.typing_compat import ParamSpec + + P = ParamSpec("P") + T = TypeVar("T") + + # # methodtools doesn't have typestubs, so give a stub + def lru_cache(maxsize: int | None = 128): + def wrapper(f): + return f + + return wrapper +else: + from methodtools import lru_cache log = structlog.get_logger(logger_name=__name__) @@ -163,13 +175,13 @@ def __init__(self, *, base_url: str | None, dry_run: bool = False, token: str, * # methods on one object prefixed with the object type (`.task_instances.update` rather than # `task_instance_update` etc.) - @methodtools.lru_cache() # type: ignore[misc] + @lru_cache() # type: ignore[misc] @property def task_instances(self) -> TaskInstanceOperations: """Operations related to TaskInstances.""" return TaskInstanceOperations(self) - @methodtools.lru_cache() # type: ignore[misc] + @lru_cache() # type: ignore[misc] @property def connections(self) -> ConnectionOperations: """Operations related to TaskInstances."""