Apply PEP-563 (Postponed Evaluation of Annotations) to core airflow (#…

…26290) This PR applies PEP-563 to all core airflow Python files - i.e. those that we want to cherry-pick to v2-* branches in order to make cherry-picking for the upcoming 2.4* releases easier. There was a separate PR (#26289) with non-core changes that are not going to be cherry-picked.. This PR is a result of combining some of the 47 PRs reviewed and approved separately (otherwise it would have been unreviewable) The history of those PRs can be changed in: https://github.com/apache/airflow/pulls?q=is%3Apr+label%3Afuture-annotations+is%3Aopen Relevant discussion: https://lists.apache.org/thread/81fr042s5d3v17v83bpo24tnrr2pp0fp Lazy consensus call: https://lists.apache.org/thread/l74nvjh8tgbtojllhwkcn7f8mfnlz4jq
apache · Sep 14, 2022 · d67ac59 · d67ac59
1 parent 06acf40
commit d67ac59
Show file tree

Hide file tree

Showing 1,097 changed files with 4,444 additions and 3,462 deletions.
diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml
@@ -215,7 +215,11 @@ repos:
  args: ["--py37-plus"]
  # We need to exclude gcs hook from pyupgrade because it has public "list" command which clashes
  # with `list` that is used as type
- exclude: ^airflow/_vendor/|^airflow/providers/google/cloud/hooks/gcs.py$
+ # Test Python tests if different kinds of typing including one that does not have
+ # __future__ annotations, so it should be left without automated upgrade
+ # BaseOperator is disabled because replacing ClassVar[List[Type with corresponding list/type causes the attr to fail
+ # see https://github.com/apache/airflow/pull/26290#issuecomment-1246014807
+ exclude: ^airflow/_vendor/|^airflow/providers/google/cloud/hooks/gcs.py$|^test/decorators/test_python.py$|^airflow/models/baseoperator.py$
  - repo: https://github.com/pre-commit/pygrep-hooks
  rev: v1.9.0
  hooks:

diff --git a/Dockerfile.ci b/Dockerfile.ci
@@ -908,7 +908,12 @@ else
  SELECTED_TESTS=()
  for provider in ${BASH_REMATCH[1]//,/ }
  do
- SELECTED_TESTS+=("tests/providers/${provider//./\/}")
+ providers_dir="tests/providers/${provider//./\/}"
+ if [[ -d ${providers_dir} ]]; then
+ SELECTED_TESTS+=("${providers_dir}")
+ else
+ echo "${COLOR_YELLOW}Skip ${providers_dir} as the directory does not exist.${COLOR_RESET}"
+ fi
  done
  else
  echo

diff --git a/airflow/__init__.py b/airflow/__init__.py
@@ -15,8 +15,6 @@
 # KIND, either express or implied. See the License for the
 # specific language governing permissions and limitations
 # under the License.
-#
-
 """
 Authentication is implemented using flask_login and different environments can
 implement their own login mechanisms by providing an `airflow_login` module
@@ -25,11 +23,10 @@
 
 isort:skip_file
 """
+from __future__ import annotations
 
 # flake8: noqa: F401
 
-from __future__ import annotations
-
 import os
 import sys
 from typing import Callable

diff --git a/airflow/__main__.py b/airflow/__main__.py
@@ -17,8 +17,9 @@
 # KIND, either express or implied. See the License for the
 # specific language governing permissions and limitations
 # under the License.
-
 """Main executable module"""
+from __future__ import annotations
+
 import os
 
 import argcomplete

diff --git a/airflow/api/__init__.py b/airflow/api/__init__.py
@@ -16,6 +16,8 @@
 # specific language governing permissions and limitations
 # under the License.
 """Authentication backend"""
+from __future__ import annotations
+
 import logging
 from importlib import import_module
 

diff --git a/airflow/api/auth/backend/basic_auth.py b/airflow/api/auth/backend/basic_auth.py
@@ -15,8 +15,10 @@
 # specific language governing permissions and limitations
 # under the License.
 """Basic authentication backend"""
+from __future__ import annotations
+
 from functools import wraps
-from typing import Any, Callable, Optional, Tuple, TypeVar, Union, cast
+from typing import Any, Callable, TypeVar, cast
 
 from flask import Response, request
 from flask_appbuilder.const import AUTH_LDAP
@@ -25,7 +27,7 @@
 from airflow.utils.airflow_flask_app import get_airflow_app
 from airflow.www.fab_security.sqla.models import User
 
-CLIENT_AUTH: Optional[Union[Tuple[str, str], Any]] = None
+CLIENT_AUTH: tuple[str, str] | Any | None = None
 
 
 def init_app(_):
@@ -35,7 +37,7 @@ def init_app(_):
 T = TypeVar("T", bound=Callable)
 
 
-def auth_current_user() -> Optional[User]:
+def auth_current_user() -> User | None:
  """Authenticate and set current user if Authorization header exists"""
  auth = request.authorization
  if auth is None or not auth.username or not auth.password:

diff --git a/airflow/api/auth/backend/default.py b/airflow/api/auth/backend/default.py
@@ -16,10 +16,12 @@
 # specific language governing permissions and limitations
 # under the License.
 """Default authentication backend - everything is allowed"""
+from __future__ import annotations
+
 from functools import wraps
-from typing import Any, Callable, Optional, Tuple, TypeVar, Union, cast
+from typing import Any, Callable, TypeVar, cast
 
-CLIENT_AUTH: Optional[Union[Tuple[str, str], Any]] = None
+CLIENT_AUTH: tuple[str, str] | Any | None = None
 
 
 def init_app(_):

diff --git a/airflow/api/auth/backend/deny_all.py b/airflow/api/auth/backend/deny_all.py
@@ -16,12 +16,14 @@
 # specific language governing permissions and limitations
 # under the License.
 """Authentication backend that denies all requests"""
+from __future__ import annotations
+
 from functools import wraps
-from typing import Any, Callable, Optional, Tuple, TypeVar, Union, cast
+from typing import Any, Callable, TypeVar, cast
 
 from flask import Response
 
-CLIENT_AUTH: Optional[Union[Tuple[str, str], Any]] = None
+CLIENT_AUTH: tuple[str, str] | Any | None = None
 
 
 def init_app(_):

diff --git a/airflow/api/auth/backend/kerberos_auth.py b/airflow/api/auth/backend/kerberos_auth.py
@@ -14,6 +14,7 @@
 # KIND, either express or implied. See the License for the
 # specific language governing permissions and limitations
 # under the License.
+from __future__ import annotations
 
 #
 # Copyright (c) 2013, Michael Komitee
@@ -43,7 +44,7 @@
 import logging
 import os
 from functools import wraps
-from typing import Any, Callable, Optional, Tuple, TypeVar, Union, cast
+from typing import Any, Callable, TypeVar, cast
 
 import kerberos
 from flask import Response, _request_ctx_stack as stack, g, make_response, request # type: ignore
@@ -55,7 +56,7 @@
 log = logging.getLogger(__name__)
 
 
-CLIENT_AUTH: Optional[Union[Tuple[str, str], Any]] = HTTPKerberosAuth(service='airflow')
+CLIENT_AUTH: tuple[str, str] | Any | None = HTTPKerberosAuth(service='airflow')
 
 
 class KerberosService:

diff --git a/airflow/api/auth/backend/session.py b/airflow/api/auth/backend/session.py
@@ -15,12 +15,14 @@
 # specific language governing permissions and limitations
 # under the License.
 """Session authentication backend"""
+from __future__ import annotations
+
 from functools import wraps
-from typing import Any, Callable, Optional, Tuple, TypeVar, Union, cast
+from typing import Any, Callable, TypeVar, cast
 
 from flask import Response, g
 
-CLIENT_AUTH: Optional[Union[Tuple[str, str], Any]] = None
+CLIENT_AUTH: tuple[str, str] | Any | None = None
 
 
 def init_app(_):

diff --git a/airflow/api/client/__init__.py b/airflow/api/client/__init__.py
@@ -16,6 +16,8 @@
 # specific language governing permissions and limitations
 # under the License.
 """API Client that allows interacting with Airflow API"""
+from __future__ import annotations
+
 from importlib import import_module
 from typing import Any
 

diff --git a/airflow/api/client/api_client.py b/airflow/api/client/api_client.py
@@ -16,6 +16,8 @@
 # specific language governing permissions and limitations
 # under the License.
 """Client for all the API clients."""
+from __future__ import annotations
+
 import httpx
 
 

diff --git a/airflow/api/client/json_client.py b/airflow/api/client/json_client.py
@@ -16,6 +16,7 @@
 # specific language governing permissions and limitations
 # under the License.
 """JSON API Client"""
+from __future__ import annotations
 
 from urllib.parse import urljoin
 

diff --git a/airflow/api/client/local_client.py b/airflow/api/client/local_client.py
@@ -16,6 +16,7 @@
 # specific language governing permissions and limitations
 # under the License.
 """Local client API"""
+from __future__ import annotations
 
 from airflow.api.client import api_client
 from airflow.api.common import delete_dag, trigger_dag

diff --git a/airflow/api/common/delete_dag.py b/airflow/api/common/delete_dag.py
@@ -16,6 +16,8 @@
 # specific language governing permissions and limitations
 # under the License.
 """Delete DAGs APIs."""
+from __future__ import annotations
+
 import logging
 
 from sqlalchemy import and_, or_

diff --git a/airflow/api/common/experimental/__init__.py b/airflow/api/common/experimental/__init__.py
@@ -16,14 +16,15 @@
 # specific language governing permissions and limitations
 # under the License.
 """Experimental APIs."""
+from __future__ import annotations
+
 from datetime import datetime
-from typing import Optional
 
 from airflow.exceptions import DagNotFound, DagRunNotFound, TaskNotFound
 from airflow.models import DagBag, DagModel, DagRun
 
 
-def check_and_get_dag(dag_id: str, task_id: Optional[str] = None) -> DagModel:
+def check_and_get_dag(dag_id: str, task_id: str | None = None) -> DagModel:
  """Checks that DAG exists and in case it is specified that Task exist"""
  dag_model = DagModel.get_current(dag_id)
  if dag_model is None:

diff --git a/airflow/api/common/experimental/delete_dag.py b/airflow/api/common/experimental/delete_dag.py
@@ -15,6 +15,8 @@
 # KIND, either express or implied. See the License for the
 # specific language governing permissions and limitations
 # under the License.
+from __future__ import annotations
+
 import warnings
 
 from airflow.api.common.delete_dag import * # noqa

diff --git a/airflow/api/common/experimental/get_code.py b/airflow/api/common/experimental/get_code.py
@@ -16,6 +16,8 @@
 # specific language governing permissions and limitations
 # under the License.
 """Get code APIs."""
+from __future__ import annotations
+
 from deprecated import deprecated
 
 from airflow.api.common.experimental import check_and_get_dag

diff --git a/airflow/api/common/experimental/get_dag_run_state.py b/airflow/api/common/experimental/get_dag_run_state.py
@@ -16,16 +16,17 @@
 # specific language governing permissions and limitations
 # under the License.
 """DAG run APIs."""
+from __future__ import annotations
+
 from datetime import datetime
-from typing import Dict
 
 from deprecated import deprecated
 
 from airflow.api.common.experimental import check_and_get_dag, check_and_get_dagrun
 
 
 @deprecated(reason="Use DagRun().get_state() instead", version="2.2.4")
-def get_dag_run_state(dag_id: str, execution_date: datetime) -> Dict[str, str]:
+def get_dag_run_state(dag_id: str, execution_date: datetime) -> dict[str, str]:
  """Return the Dag Run state identified by the given dag_id and execution_date.
 
  :param dag_id: DAG id

diff --git a/airflow/api/common/experimental/get_dag_runs.py b/airflow/api/common/experimental/get_dag_runs.py
@@ -16,7 +16,9 @@
 # specific language governing permissions and limitations
 # under the License.
 """DAG runs APIs."""
-from typing import Any, Dict, List, Optional
+from __future__ import annotations
+
+from typing import Any
 
 from flask import url_for
 
@@ -25,7 +27,7 @@
 from airflow.utils.state import DagRunState
 
 
-def get_dag_runs(dag_id: str, state: Optional[str] = None) -> List[Dict[str, Any]]:
+def get_dag_runs(dag_id: str, state: str | None = None) -> list[dict[str, Any]]:
  """
  Returns a list of Dag Runs for a specific DAG ID.
 

diff --git a/airflow/api/common/experimental/get_lineage.py b/airflow/api/common/experimental/get_lineage.py
@@ -16,9 +16,11 @@
 # specific language governing permissions and limitations
 # under the License.
 """Lineage apis"""
+from __future__ import annotations
+
 import collections
 import datetime
-from typing import Any, Dict
+from typing import Any
 
 from sqlalchemy.orm import Session
 
@@ -31,15 +33,15 @@
 @provide_session
 def get_lineage(
  dag_id: str, execution_date: datetime.datetime, *, session: Session = NEW_SESSION
-) -> Dict[str, Dict[str, Any]]:
+) -> dict[str, dict[str, Any]]:
  """Gets the lineage information for dag specified."""
  dag = check_and_get_dag(dag_id)
  dagrun = check_and_get_dagrun(dag, execution_date)
 
  inlets = XCom.get_many(dag_ids=dag_id, run_id=dagrun.run_id, key=PIPELINE_INLETS, session=session)
  outlets = XCom.get_many(dag_ids=dag_id, run_id=dagrun.run_id, key=PIPELINE_OUTLETS, session=session)
 
- lineage: Dict[str, Dict[str, Any]] = collections.defaultdict(dict)
+ lineage: dict[str, dict[str, Any]] = collections.defaultdict(dict)
  for meta in inlets:
  lineage[meta.task_id]["inlets"] = meta.value
  for meta in outlets:

diff --git a/airflow/api/common/experimental/get_task.py b/airflow/api/common/experimental/get_task.py
@@ -16,6 +16,8 @@
 # specific language governing permissions and limitations
 # under the License.
 """Task APIs.."""
+from __future__ import annotations
+
 from deprecated import deprecated
 
 from airflow.api.common.experimental import check_and_get_dag

diff --git a/airflow/api/common/experimental/get_task_instance.py b/airflow/api/common/experimental/get_task_instance.py
@@ -15,7 +15,8 @@
 # KIND, either express or implied. See the License for the
 # specific language governing permissions and limitations
 # under the License.
-"""Task Instance APIs."""
+from __future__ import annotations
+
 from datetime import datetime
 
 from deprecated import deprecated

diff --git a/airflow/api/common/experimental/mark_tasks.py b/airflow/api/common/experimental/mark_tasks.py
@@ -15,6 +15,9 @@
 # KIND, either express or implied. See the License for the
 # specific language governing permissions and limitations
 # under the License.
+"""Task Instance APIs."""
+from __future__ import annotations
+
 import warnings
 
 from airflow.api.common.mark_tasks import ( # noqa

diff --git a/airflow/api/common/experimental/pool.py b/airflow/api/common/experimental/pool.py
@@ -16,6 +16,8 @@
 # specific language governing permissions and limitations
 # under the License.
 """Pool APIs."""
+from __future__ import annotations
+
 from deprecated import deprecated
 
 from airflow.exceptions import AirflowBadRequest, PoolNotFound

diff --git a/airflow/api/common/experimental/trigger_dag.py b/airflow/api/common/experimental/trigger_dag.py
@@ -15,6 +15,7 @@
 # KIND, either express or implied. See the License for the
 # specific language governing permissions and limitations
 # under the License.
+from __future__ import annotations
 
 import warnings