From 2bd701a68a19bb61898e74cad8ff3a4e761119a6 Mon Sep 17 00:00:00 2001 From: Ash Berlin-Taylor Date: Mon, 18 Nov 2024 18:15:38 +0000 Subject: [PATCH] Improve IDE type hinting of methdtools.lru_cache in Task SDK client (#44152) mypy doesn't check the type of double decorators (hence why we still need the ignore misc) but pyright does, and methodtools doesn't have type hints, which lead to it not being able to tell what type of `client.task_instances` was --- task_sdk/src/airflow/sdk/api/client.py | 20 ++++++++++++++++---- 1 file changed, 16 insertions(+), 4 deletions(-) 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."""