Use methodtools.lru_cache instead of functools.lru_cache in class…

… methods (#37757)

* Use `methodtools.lru_cache` instead of `functools.lru_cache` in class methods

* Rename pre-commit script

.pre-commit-config.yaml

@@ -985,6 +985,13 @@ repos:
         files: \.py$
         exclude: ^.*/.*_vendor/
         additional_dependencies: ['rich>=12.4.4']
+      - id: check-compat-cache-on-methods
+        name: Check that compat cache do not use on class methods
+        entry: ./scripts/ci/pre_commit/compat_cache_on_methods.py
+        language: python
+        pass_filenames: true
+        files: ^airflow/.*\.py$
+        exclude: ^.*/.*_vendor/
       - id: lint-chart-schema
         name: Lint chart/values.schema.json file
         entry: ./scripts/ci/pre_commit/chart_schema.py

airflow/models/abstractoperator.py

@@ -23,9 +23,9 @@
 from functools import cached_property
 from typing import TYPE_CHECKING, Any, Callable, ClassVar, Collection, Iterable, Iterator, Sequence
 
+import methodtools
 from sqlalchemy import select
 
-from airflow.compat.functools import cache
 from airflow.configuration import conf
 from airflow.exceptions import AirflowException
 from airflow.models.expandinput import NotFullyPopulated
@@ -493,7 +493,7 @@ def get_extra_links(self, ti: TaskInstance, link_name: str) -> str | None:
             return link.get_link(self.unmap(None), ti.dag_run.logical_date)  # type: ignore[misc]
         return link.get_link(self.unmap(None), ti_key=ti.key)
 
-    @cache
+    @methodtools.lru_cache(maxsize=None)
     def get_parse_time_mapped_ti_count(self) -> int:
         """
         Return the number of mapped task instances that can be created on DAG run creation.

airflow/models/mappedoperator.py

@@ -24,8 +24,8 @@
 from typing import TYPE_CHECKING, Any, ClassVar, Collection, Iterable, Iterator, Mapping, Sequence, Union
 
 import attr
+import methodtools
 
-from airflow.compat.functools import cache
 from airflow.exceptions import AirflowException, UnmappableOperator
 from airflow.models.abstractoperator import (
     DEFAULT_EXECUTOR,
@@ -335,8 +335,8 @@ def __attrs_post_init__(self):
                 f"{self.task_id!r}."
             )
 
+    @methodtools.lru_cache(maxsize=None)
     @classmethod
-    @cache
     def get_serialized_fields(cls):
         # Not using 'cls' here since we only want to serialize base fields.
         return frozenset(attr.fields_dict(MappedOperator)) - {
@@ -352,8 +352,8 @@ def get_serialized_fields(cls):
             "_on_failure_fail_dagrun",
         }
 
+    @methodtools.lru_cache(maxsize=None)
     @staticmethod
-    @cache
     def deps_for(operator_class: type[BaseOperator]) -> frozenset[BaseTIDep]:
         operator_deps = operator_class.deps
         if not isinstance(operator_deps, collections.abc.Set):
@@ -789,7 +789,7 @@ def iter_mapped_dependencies(self) -> Iterator[Operator]:
         for operator, _ in XComArg.iter_xcom_references(self._get_specified_expand_input()):
             yield operator
 
-    @cache
+    @methodtools.lru_cache(maxsize=None)
     def get_parse_time_mapped_ti_count(self) -> int:
         current_count = self._get_specified_expand_input().get_parse_time_mapped_ti_count()
         try:

airflow/utils/task_group.py

@@ -25,9 +25,9 @@
 import weakref
 from typing import TYPE_CHECKING, Any, Generator, Iterator, Sequence
 
+import methodtools
 import re2
 
-from airflow.compat.functools import cache
 from airflow.exceptions import (
     AirflowDagCycleException,
     AirflowException,
@@ -586,7 +586,7 @@ def iter_mapped_dependencies(self) -> Iterator[Operator]:
         for op, _ in XComArg.iter_xcom_references(self._expand_input):
             yield op
 
-    @cache
+    @methodtools.lru_cache(maxsize=None)
     def get_parse_time_mapped_ti_count(self) -> int:
         """
         Return the Number of instances a task in this group should be mapped to, when a DAG run is created.

airflow/utils/weight_rule.py

@@ -19,7 +19,7 @@
 
 from enum import Enum
 
-from airflow.compat.functools import cache
+import methodtools
 
 
 class WeightRule(str, Enum):
@@ -34,8 +34,8 @@ def is_valid(cls, weight_rule: str) -> bool:
         """Check if weight rule is valid."""
         return weight_rule in cls.all_weight_rules()
 
+    @methodtools.lru_cache(maxsize=None)
     @classmethod
-    @cache
     def all_weight_rules(cls) -> set[str]:
         """Return all weight rules."""
         return set(cls.__members__.values())

contributing-docs/08_static_code_checks.rst

@@ -146,6 +146,8 @@ require Breeze Docker image to be built locally.
 +-----------------------------------------------------------+--------------------------------------------------------------+---------+
 | check-cncf-k8s-only-for-executors                         | Check cncf.kubernetes imports used for executors only        |         |
 +-----------------------------------------------------------+--------------------------------------------------------------+---------+
+| check-compat-cache-on-methods                             | Check that compat cache do not use on class methods          |         |
++-----------------------------------------------------------+--------------------------------------------------------------+---------+
 | check-core-deprecation-classes                            | Verify usage of Airflow deprecation classes in core          |         |
 +-----------------------------------------------------------+--------------------------------------------------------------+---------+
 | check-daysago-import-from-utils                           | Make sure days_ago is imported from airflow.utils.dates      |         |