From 72521e457ea7d342fcf41e534820a69bd8d12085 Mon Sep 17 00:00:00 2001 From: Tzu-ping Chung Date: Tue, 25 May 2021 21:36:00 +0800 Subject: [PATCH] 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. (cherry picked from commit 3db347edcfe444a67e59e8cf0019e80a02dbceab) (cherry picked from commit 1aa1bbd73e62d491becedf5933452a53fddc297a) --- airflow/cli/commands/connection_command.py | 3 +- airflow/cli/commands/kubernetes_command.py | 3 +- airflow/cli/simple_table.py | 2 +- airflow/compat/__init__.py | 16 +++++++++++ airflow/compat/functools.py | 33 ++++++++++++++++++++++ airflow/configuration.py | 3 +- airflow/kubernetes/pod_generator.py | 2 +- airflow/kubernetes/refresh_config.py | 2 +- airflow/models/baseoperator.py | 6 +--- airflow/operators/bash.py | 6 +--- airflow/operators/sql.py | 6 +--- airflow/providers_manager.py | 2 +- airflow/secrets/local_filesystem.py | 2 +- airflow/utils/log/log_reader.py | 6 +--- airflow/utils/log/secrets_masker.py | 15 +--------- airflow/utils/yaml.py | 2 +- airflow/www/views.py | 3 +- 17 files changed, 64 insertions(+), 48 deletions(-) create mode 100644 airflow/compat/__init__.py create mode 100644 airflow/compat/functools.py diff --git a/airflow/cli/commands/connection_command.py b/airflow/cli/commands/connection_command.py index 19912b6e8176..c79ba6a4d1b3 100644 --- a/airflow/cli/commands/connection_command.py +++ b/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 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 diff --git a/airflow/cli/commands/kubernetes_command.py b/airflow/cli/commands/kubernetes_command.py index daf11a3e7097..3aad61622c2e 100644 --- a/airflow/cli/commands/kubernetes_command.py +++ b/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 diff --git a/airflow/cli/simple_table.py b/airflow/cli/simple_table.py index 65e846e347e9..d17f948d2907 100644 --- a/airflow/cli/simple_table.py +++ b/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 diff --git a/airflow/compat/__init__.py b/airflow/compat/__init__.py new file mode 100644 index 000000000000..13a83393a912 --- /dev/null +++ b/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. diff --git a/airflow/compat/functools.py b/airflow/compat/functools.py new file mode 100644 index 000000000000..10b4085d2db0 --- /dev/null +++ b/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"] diff --git a/airflow/configuration.py b/airflow/configuration.py index 263ba4b62c61..ed38efdb106e 100644 --- a/airflow/configuration.py +++ b/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() -> List[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) diff --git a/airflow/kubernetes/pod_generator.py b/airflow/kubernetes/pod_generator.py index 80602e3cb9a8..4b4028f18ce0 100644 --- a/airflow/kubernetes/pod_generator.py +++ b/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 diff --git a/airflow/kubernetes/refresh_config.py b/airflow/kubernetes/refresh_config.py index 2738b8657cff..a039e7d2ff2f 100644 --- a/airflow/kubernetes/refresh_config.py +++ b/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: diff --git a/airflow/models/baseoperator.py b/airflow/models/baseoperator.py index f74c5f907a23..7af23d3ed26f 100644 --- a/airflow/models/baseoperator.py +++ b/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 diff --git a/airflow/operators/bash.py b/airflow/operators/bash.py index 1cc85e1fb24d..66a70a957b3c 100644 --- a/airflow/operators/bash.py +++ b/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 diff --git a/airflow/operators/sql.py b/airflow/operators/sql.py index 769cf26dc948..8347bcb41504 100644 --- a/airflow/operators/sql.py +++ b/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 diff --git a/airflow/providers_manager.py b/airflow/providers_manager.py index 420135b249ab..9f1469c8ae53 100644 --- a/airflow/providers_manager.py +++ b/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: diff --git a/airflow/secrets/local_filesystem.py b/airflow/secrets/local_filesystem.py index 3ec20e12d41f..d23969fbecb5 100644 --- a/airflow/secrets/local_filesystem.py +++ b/airflow/secrets/local_filesystem.py @@ -25,7 +25,6 @@ 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, @@ -33,6 +32,7 @@ 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 diff --git a/airflow/utils/log/log_reader.py b/airflow/utils/log/log_reader.py index f4096ebaeada..0e1f69198d01 100644 --- a/airflow/utils/log/log_reader.py +++ b/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 diff --git a/airflow/utils/log/secrets_masker.py b/airflow/utils/log/secrets_masker.py index 1796cbc331aa..8bb7557a9df2 100644 --- a/airflow/utils/log/secrets_masker.py +++ b/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 diff --git a/airflow/utils/yaml.py b/airflow/utils/yaml.py index e3be61c31501..c4526313ddf6 100644 --- a/airflow/utils/yaml.py +++ b/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: diff --git a/airflow/www/views.py b/airflow/www/views.py index 350119954062..b66374944345 100644 --- a/airflow/www/views.py +++ b/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