Introduce compat shim airflow.compat.functools (#15969)

This module shims 'cached_property' and 'cache' so modules don't need to
all do their own ad-hoc try-except ImportError.

airflow/cli/commands/connection_command.py

@@ -24,13 +24,12 @@
 
 from sqlalchemy.orm import exc
 
-import airflow.utils.yaml as yaml
 from airflow.cli.simple_table import AirflowConsole
 from airflow.exceptions import AirflowNotFoundException
 from airflow.hooks.base import BaseHook
 from airflow.models import Connection
 from airflow.secrets.local_filesystem import _create_connection, load_connections_dict
-from airflow.utils import cli as cli_utils
+from airflow.utils import cli as cli_utils, yaml
 from airflow.utils.cli import suppress_logs_and_warning
 from airflow.utils.session import create_session
 

airflow/cli/commands/kubernetes_command.py

@@ -22,14 +22,13 @@
 from kubernetes.client.api_client import ApiClient
 from kubernetes.client.rest import ApiException
 
-import airflow.utils.yaml as yaml
 from airflow.executors.kubernetes_executor import KubeConfig, create_pod_id
 from airflow.kubernetes import pod_generator
 from airflow.kubernetes.kube_client import get_kube_client
 from airflow.kubernetes.pod_generator import PodGenerator
 from airflow.models import TaskInstance
 from airflow.settings import pod_mutation_hook
-from airflow.utils import cli as cli_utils
+from airflow.utils import cli as cli_utils, yaml
 from airflow.utils.cli import get_dag
 
 

airflow/cli/simple_table.py

@@ -24,8 +24,8 @@
 from rich.table import Table
 from tabulate import tabulate
 
-import airflow.utils.yaml as yaml
 from airflow.plugins_manager import PluginsDirectorySource
+from airflow.utils import yaml
 from airflow.utils.platform import is_tty
 
 

airflow/compat/__init__.py

@@ -0,0 +1,16 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.

airflow/compat/functools.py

@@ -0,0 +1,33 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+import sys
+
+if sys.version_info >= (3, 8):
+    from functools import cached_property  # pylint: disable=no-name-in-module
+else:
+    from cached_property import cached_property
+
+if sys.version_info >= (3, 9):
+    from functools import cache  # pylint: disable=no-name-in-module
+else:
+    from functools import lru_cache
+
+    cache = lru_cache(maxsize=None)
+
+
+__all__ = ["cache", "cached_property"]

airflow/configuration.py

@@ -36,6 +36,7 @@
 
 from airflow.exceptions import AirflowConfigException
 from airflow.secrets import DEFAULT_SECRETS_SEARCH_PATH, BaseSecretsBackend
+from airflow.utils import yaml
 from airflow.utils.module_loading import import_string
 
 log = logging.getLogger(__name__)
@@ -97,8 +98,6 @@ def default_config_yaml() -> dict:
 
     :return: Python dictionary containing configs & their info
     """
-    import airflow.utils.yaml as yaml
-
     with open(_default_config_file_path('config.yml')) as config_file:
         return yaml.safe_load(config_file)
 

airflow/kubernetes/pod_generator.py

@@ -34,9 +34,9 @@
 from kubernetes.client import models as k8s
 from kubernetes.client.api_client import ApiClient
 
-import airflow.utils.yaml as yaml
 from airflow.exceptions import AirflowConfigException
 from airflow.kubernetes.pod_generator_deprecated import PodDefaults, PodGenerator as PodGeneratorDeprecated
+from airflow.utils import yaml
 from airflow.version import version as airflow_version
 
 MAX_LABEL_LEN = 63

airflow/kubernetes/refresh_config.py

@@ -31,7 +31,7 @@
 from kubernetes.config.exec_provider import ExecProvider
 from kubernetes.config.kube_config import KUBE_CONFIG_DEFAULT_LOCATION, KubeConfigLoader
 
-import airflow.utils.yaml as yaml
+from airflow.utils import yaml
 
 
 def _parse_timestamp(ts_str: str) -> int:

airflow/models/baseoperator.py

@@ -46,15 +46,11 @@
 
 import attr
 import jinja2
-
-try:
-    from functools import cached_property
-except ImportError:
-    from cached_property import cached_property
 from dateutil.relativedelta import relativedelta
 from sqlalchemy.orm import Session
 
 import airflow.templates
+from airflow.compat.functools import cached_property
 from airflow.configuration import conf
 from airflow.exceptions import AirflowException
 from airflow.lineage import apply_lineage, prepare_lineage

airflow/operators/bash.py

@@ -18,11 +18,7 @@
 import os
 from typing import Dict, Optional
 
-try:
-    from functools import cached_property
-except ImportError:
-    from cached_property import cached_property
-
+from airflow.compat.functools import cached_property
 from airflow.exceptions import AirflowException, AirflowSkipException
 from airflow.hooks.subprocess import SubprocessHook
 from airflow.models import BaseOperator

airflow/operators/sql.py

@@ -18,11 +18,7 @@
 from distutils.util import strtobool
 from typing import Any, Dict, Iterable, List, Mapping, Optional, SupportsAbs, Union
 
-try:
-    from functools import cached_property
-except ImportError:
-    from cached_property import cached_property
-
+from airflow.compat.functools import cached_property
 from airflow.exceptions import AirflowException
 from airflow.hooks.base import BaseHook
 from airflow.hooks.dbapi import DbApiHook

airflow/providers_manager.py

@@ -27,7 +27,7 @@
 import jsonschema
 from wtforms import Field
 
-import airflow.utils.yaml as yaml
+from airflow.utils import yaml
 from airflow.utils.entry_points import entry_points_with_dist
 
 try:

airflow/secrets/local_filesystem.py

@@ -25,14 +25,14 @@
 from json import JSONDecodeError
 from typing import TYPE_CHECKING, Any, Dict, List, Optional, Set, Tuple
 
-import airflow.utils.yaml as yaml
 from airflow.exceptions import (
     AirflowException,
     AirflowFileParseException,
     ConnectionNotUnique,
     FileSyntaxError,
 )
 from airflow.secrets.base_secrets import BaseSecretsBackend
+from airflow.utils import yaml
 from airflow.utils.file import COMMENT_PATTERN
 from airflow.utils.log.logging_mixin import LoggingMixin
 

airflow/utils/log/log_reader.py

@@ -18,11 +18,7 @@
 import logging
 from typing import Dict, Iterator, List, Optional, Tuple
 
-try:
-    from functools import cached_property
-except ImportError:
-    from cached_property import cached_property
-
+from airflow.compat.functools import cached_property
 from airflow.configuration import conf
 from airflow.models import TaskInstance
 from airflow.utils.helpers import render_log_filename

airflow/utils/log/secrets_masker.py

@@ -20,20 +20,7 @@
 import re
 from typing import TYPE_CHECKING, Iterable, Optional, Set, TypeVar, Union
 
-try:
-    # 3.8+
-    from functools import cached_property
-except ImportError:
-    from cached_property import cached_property
-
-try:
-    # 3.9+
-    from functools import cache
-except ImportError:
-    from functools import lru_cache
-
-    cache = lru_cache(maxsize=None)
-
+from airflow.compat.functools import cache, cached_property
 
 if TYPE_CHECKING:
     from airflow.typing_compat import RePatternType

airflow/utils/yaml.py

@@ -30,7 +30,7 @@
 from typing import TYPE_CHECKING, Any, BinaryIO, TextIO, Union, cast
 
 if TYPE_CHECKING:
-    from yaml.error import MarkedYAMLError  # noqa
+    from yaml.error import MarkedYAMLError, YAMLError  # noqa
 
 
 def safe_load(stream: Union[bytes, str, BinaryIO, TextIO]) -> Any:

airflow/www/views.py

@@ -82,7 +82,6 @@
 from wtforms.validators import InputRequired
 
 import airflow
-import airflow.utils.yaml as yaml
 from airflow import models, plugins_manager, settings
 from airflow.api.common.experimental.mark_tasks import (
     set_dag_run_state_to_failed,
@@ -103,7 +102,7 @@
 from airflow.security import permissions
 from airflow.ti_deps.dep_context import DepContext
 from airflow.ti_deps.dependencies_deps import RUNNING_DEPS, SCHEDULER_QUEUED_DEPS
-from airflow.utils import json as utils_json, timezone
+from airflow.utils import json as utils_json, timezone, yaml
 from airflow.utils.dates import infer_time_unit, scale_time_units
 from airflow.utils.docs import get_docs_url
 from airflow.utils.helpers import alchemy_to_dict