From 7fec8221434322f4502fce9e4629bbbc5d4520d1 Mon Sep 17 00:00:00 2001 From: Andrey Anshin Date: Sun, 27 Aug 2023 13:03:31 +0400 Subject: [PATCH] Always use `Literal` from `typing_extensions` --- airflow/cli/commands/task_command.py | 4 ++-- airflow/models/dag.py | 2 +- airflow/models/dagrun.py | 2 +- airflow/models/mappedoperator.py | 2 +- airflow/models/taskinstance.py | 3 ++- airflow/providers/amazon/aws/sensors/sqs.py | 3 ++- airflow/providers/amazon/aws/transfers/sql_to_s3.py | 4 +++- airflow/providers/amazon/aws/triggers/sqs.py | 4 +++- airflow/providers/amazon/aws/utils/sqs.py | 3 ++- airflow/providers/cncf/kubernetes/operators/pod.py | 2 +- airflow/providers/cncf/kubernetes/utils/pod_manager.py | 4 ++-- airflow/providers/mongo/hooks/mongo.py | 2 +- airflow/providers/salesforce/operators/bulk.py | 3 ++- airflow/providers_manager.py | 2 +- airflow/typing_compat.py | 8 +++----- airflow/utils/dates.py | 2 +- dev/breeze/src/airflow_breeze/utils/selective_checks.py | 5 +---- 17 files changed, 29 insertions(+), 26 deletions(-) diff --git a/airflow/cli/commands/task_command.py b/airflow/cli/commands/task_command.py index 97b290f0780fb..8884f050fe27c 100644 --- a/airflow/cli/commands/task_command.py +++ b/airflow/cli/commands/task_command.py @@ -25,13 +25,14 @@ import sys import textwrap from contextlib import contextmanager, redirect_stderr, redirect_stdout, suppress -from typing import Generator, Union, cast +from typing import Generator, Protocol, Union, cast import pendulum from pendulum.parsing.exceptions import ParserError from sqlalchemy import select from sqlalchemy.orm.exc import NoResultFound from sqlalchemy.orm.session import Session +from typing_extensions import Literal from airflow import settings from airflow.cli.simple_table import AirflowConsole @@ -50,7 +51,6 @@ from airflow.settings import IS_K8S_EXECUTOR_POD from airflow.ti_deps.dep_context import DepContext from airflow.ti_deps.dependencies_deps import SCHEDULER_QUEUED_DEPS -from airflow.typing_compat import Literal, Protocol from airflow.utils import cli as cli_utils from airflow.utils.cli import ( get_dag, diff --git a/airflow/models/dag.py b/airflow/models/dag.py index 930e6b2cd6e63..5979c51c518fa 100644 --- a/airflow/models/dag.py +++ b/airflow/models/dag.py @@ -131,13 +131,13 @@ from pendulum.tz.timezone import Timezone from sqlalchemy.orm.query import Query from sqlalchemy.orm.session import Session + from typing_extensions import Literal from airflow.datasets import Dataset from airflow.decorators import TaskDecoratorCollection from airflow.models.dagbag import DagBag from airflow.models.operator import Operator from airflow.models.slamiss import SlaMiss - from airflow.typing_compat import Literal from airflow.utils.task_group import TaskGroup log = logging.getLogger(__name__) diff --git a/airflow/models/dagrun.py b/airflow/models/dagrun.py index 43ed97e518835..c528b6114bff5 100644 --- a/airflow/models/dagrun.py +++ b/airflow/models/dagrun.py @@ -72,10 +72,10 @@ from datetime import datetime from sqlalchemy.orm import Query, Session + from typing_extensions import Literal from airflow.models.dag import DAG from airflow.models.operator import Operator - from airflow.typing_compat import Literal from airflow.utils.types import ArgNotSet CreatedTasks = TypeVar("CreatedTasks", Iterator["dict[str, Any]"], Iterator[TI]) diff --git a/airflow/models/mappedoperator.py b/airflow/models/mappedoperator.py index 37882dcd78dfb..09785aaaf52db 100644 --- a/airflow/models/mappedoperator.py +++ b/airflow/models/mappedoperator.py @@ -25,6 +25,7 @@ from typing import TYPE_CHECKING, Any, ClassVar, Collection, Iterable, Iterator, Mapping, Sequence, Union import attr +from typing_extensions import Literal from airflow import settings from airflow.compat.functools import cache @@ -51,7 +52,6 @@ from airflow.models.pool import Pool from airflow.serialization.enums import DagAttributeTypes from airflow.ti_deps.deps.mapped_task_expanded import MappedTaskIsExpanded -from airflow.typing_compat import Literal from airflow.utils.context import context_update_for_unmapped from airflow.utils.helpers import is_container, prevent_duplicates from airflow.utils.types import NOTSET diff --git a/airflow/models/taskinstance.py b/airflow/models/taskinstance.py index ff1b190a5eab8..b28d59bee088e 100644 --- a/airflow/models/taskinstance.py +++ b/airflow/models/taskinstance.py @@ -135,6 +135,7 @@ from sqlalchemy.orm.session import Session from sqlalchemy.sql.elements import BooleanClauseList from sqlalchemy.sql.expression import ColumnOperators + from typing_extensions import Literal from airflow.models.abstractoperator import TaskStateChangeCallback from airflow.models.baseoperator import BaseOperator @@ -143,7 +144,7 @@ from airflow.models.dataset import DatasetEvent from airflow.models.operator import Operator from airflow.timetables.base import DataInterval - from airflow.typing_compat import Literal, TypeGuard + from airflow.typing_compat import TypeGuard from airflow.utils.task_group import TaskGroup # This is a workaround because mypy doesn't work with hybrid_property diff --git a/airflow/providers/amazon/aws/sensors/sqs.py b/airflow/providers/amazon/aws/sensors/sqs.py index c4bf2621f3d9f..43191d1445e10 100644 --- a/airflow/providers/amazon/aws/sensors/sqs.py +++ b/airflow/providers/amazon/aws/sensors/sqs.py @@ -19,9 +19,10 @@ from __future__ import annotations from functools import cached_property -from typing import TYPE_CHECKING, Any, Collection, Literal, Sequence +from typing import TYPE_CHECKING, Any, Collection, Sequence from deprecated import deprecated +from typing_extensions import Literal from airflow.configuration import conf from airflow.exceptions import AirflowException diff --git a/airflow/providers/amazon/aws/transfers/sql_to_s3.py b/airflow/providers/amazon/aws/transfers/sql_to_s3.py index 1302927bfd5fc..e0860620234bc 100644 --- a/airflow/providers/amazon/aws/transfers/sql_to_s3.py +++ b/airflow/providers/amazon/aws/transfers/sql_to_s3.py @@ -20,7 +20,9 @@ import enum from collections import namedtuple from tempfile import NamedTemporaryFile -from typing import TYPE_CHECKING, Iterable, Literal, Mapping, Sequence +from typing import TYPE_CHECKING, Iterable, Mapping, Sequence + +from typing_extensions import Literal from airflow.exceptions import AirflowException from airflow.hooks.base import BaseHook diff --git a/airflow/providers/amazon/aws/triggers/sqs.py b/airflow/providers/amazon/aws/triggers/sqs.py index 68f85c5f53019..7b3b0f312b34c 100644 --- a/airflow/providers/amazon/aws/triggers/sqs.py +++ b/airflow/providers/amazon/aws/triggers/sqs.py @@ -17,7 +17,9 @@ from __future__ import annotations import asyncio -from typing import Any, AsyncIterator, Collection, Literal +from typing import Any, AsyncIterator, Collection + +from typing_extensions import Literal from airflow.exceptions import AirflowException from airflow.providers.amazon.aws.hooks.base_aws import BaseAwsConnection diff --git a/airflow/providers/amazon/aws/utils/sqs.py b/airflow/providers/amazon/aws/utils/sqs.py index 2b081e5259464..ea0c7afea10eb 100644 --- a/airflow/providers/amazon/aws/utils/sqs.py +++ b/airflow/providers/amazon/aws/utils/sqs.py @@ -18,9 +18,10 @@ import json import logging -from typing import Any, Literal +from typing import Any from jsonpath_ng import parse +from typing_extensions import Literal log = logging.getLogger(__name__) diff --git a/airflow/providers/cncf/kubernetes/operators/pod.py b/airflow/providers/cncf/kubernetes/operators/pod.py index f4fe0b3b2ac0b..2a02a0bbf6a4f 100644 --- a/airflow/providers/cncf/kubernetes/operators/pod.py +++ b/airflow/providers/cncf/kubernetes/operators/pod.py @@ -71,9 +71,9 @@ if TYPE_CHECKING: import jinja2 + from typing_extensions import Literal from airflow.providers.cncf.kubernetes.secret import Secret - from airflow.typing_compat import Literal from airflow.utils.context import Context alphanum_lower = string.ascii_lowercase + string.digits diff --git a/airflow/providers/cncf/kubernetes/utils/pod_manager.py b/airflow/providers/cncf/kubernetes/utils/pod_manager.py index 06371f1f1a98a..46d1d4bcefdbf 100644 --- a/airflow/providers/cncf/kubernetes/utils/pod_manager.py +++ b/airflow/providers/cncf/kubernetes/utils/pod_manager.py @@ -28,7 +28,7 @@ from contextlib import closing, suppress from dataclasses import dataclass from datetime import timedelta -from typing import TYPE_CHECKING, Generator, cast +from typing import TYPE_CHECKING, Generator, Protocol, cast import pendulum import tenacity @@ -38,11 +38,11 @@ from pendulum import DateTime from pendulum.parsing.exceptions import ParserError from tenacity import before_log +from typing_extensions import Literal from urllib3.exceptions import HTTPError as BaseHTTPError from airflow.exceptions import AirflowException, AirflowProviderDeprecationWarning from airflow.providers.cncf.kubernetes.pod_generator import PodDefaults -from airflow.typing_compat import Literal, Protocol from airflow.utils.log.logging_mixin import LoggingMixin from airflow.utils.timezone import utcnow diff --git a/airflow/providers/mongo/hooks/mongo.py b/airflow/providers/mongo/hooks/mongo.py index b207c29c13ae2..fca855a51d98d 100644 --- a/airflow/providers/mongo/hooks/mongo.py +++ b/airflow/providers/mongo/hooks/mongo.py @@ -30,7 +30,7 @@ if TYPE_CHECKING: from types import TracebackType - from airflow.typing_compat import Literal + from typing_extensions import Literal class MongoHook(BaseHook): diff --git a/airflow/providers/salesforce/operators/bulk.py b/airflow/providers/salesforce/operators/bulk.py index 554d96656bab4..98f718cb02429 100644 --- a/airflow/providers/salesforce/operators/bulk.py +++ b/airflow/providers/salesforce/operators/bulk.py @@ -22,7 +22,8 @@ from airflow.providers.salesforce.hooks.salesforce import SalesforceHook if TYPE_CHECKING: - from airflow.typing_compat import Literal + from typing_extensions import Literal + from airflow.utils.context import Context diff --git a/airflow/providers_manager.py b/airflow/providers_manager.py index 48de7776b7379..002747c33255d 100644 --- a/airflow/providers_manager.py +++ b/airflow/providers_manager.py @@ -33,11 +33,11 @@ from typing import TYPE_CHECKING, Any, Callable, MutableMapping, NamedTuple, TypeVar, cast from packaging.utils import canonicalize_name +from typing_extensions import Literal from airflow.exceptions import AirflowOptionalProviderFeatureException from airflow.hooks.filesystem import FSHook from airflow.hooks.package_index import PackageIndexHook -from airflow.typing_compat import Literal from airflow.utils import yaml from airflow.utils.entry_points import entry_points_with_dist from airflow.utils.log.logging_mixin import LoggingMixin diff --git a/airflow/typing_compat.py b/airflow/typing_compat.py index a17c6b621ec6f..95ef4595522ed 100644 --- a/airflow/typing_compat.py +++ b/airflow/typing_compat.py @@ -30,11 +30,9 @@ import sys from typing import Protocol, TypedDict, runtime_checkable -# Literal in 3.8 is limited to one single argument, not e.g. "Literal[1, 2]". -if sys.version_info >= (3, 9): - from typing import Literal -else: - from typing import Literal +# Literal from typing module has various issues in different Python versions, see: +# https://typing-extensions.readthedocs.io/en/latest/#Literal +from typing_extensions import Literal if sys.version_info >= (3, 10): from typing import ParamSpec, TypeGuard diff --git a/airflow/utils/dates.py b/airflow/utils/dates.py index 9351c64e2a574..bcf8f0a183c3c 100644 --- a/airflow/utils/dates.py +++ b/airflow/utils/dates.py @@ -23,9 +23,9 @@ from croniter import croniter from dateutil.relativedelta import relativedelta # for doctest +from typing_extensions import Literal from airflow.exceptions import RemovedInAirflow3Warning -from airflow.typing_compat import Literal from airflow.utils import timezone cron_presets: dict[str, str] = { diff --git a/dev/breeze/src/airflow_breeze/utils/selective_checks.py b/dev/breeze/src/airflow_breeze/utils/selective_checks.py index e6b350df19a8e..a1c6e1ce1e6f8 100644 --- a/dev/breeze/src/airflow_breeze/utils/selective_checks.py +++ b/dev/breeze/src/airflow_breeze/utils/selective_checks.py @@ -24,10 +24,7 @@ from re import match from typing import Any, Dict, List, TypeVar -if sys.version_info >= (3, 9): - from typing import Literal -else: - from typing import Literal +from typing_extensions import Literal from airflow_breeze.global_constants import ( ALL_PYTHON_MAJOR_MINOR_VERSIONS,