Fix new MyPy errors in main (#22884)

Those MyPe errors are side effect of some new dependencies.
apache · Apr 10, 2022 · 6933022 · 6933022
1 parent 0ae0f7e
commit 6933022
Show file tree

Hide file tree

Showing 133 changed files with 1,178 additions and 1,021 deletions.
diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml
@@ -443,9 +443,6 @@ repos:
  pass_filenames: true
  files: >
  (?x)
- ^airflow/hooks/.*$|
- ^airflow/operators/.*$|
- ^airflow/sensors/.*$|
  ^airflow/providers/.*\.py$
  exclude: ^airflow/_vendor/
  - id: provide-create-sessions

diff --git a/airflow/cli/commands/celery_command.py b/airflow/cli/commands/celery_command.py
@@ -23,7 +23,7 @@
 import daemon
 import psutil
 import sqlalchemy.exc
-from celery import maybe_patch_concurrency
+from celery import maybe_patch_concurrency # type: ignore[attr-defined]
 from daemon.pidfile import TimeoutPIDLockFile
 from lockfile.pidlockfile import read_pid_from_pidfile, remove_existing_pidfile
 

diff --git a/airflow/models/abstractoperator.py b/airflow/models/abstractoperator.py
@@ -29,6 +29,7 @@
  Iterable,
  List,
  Optional,
+ Sequence,
  Set,
  Type,
  Union,
@@ -92,7 +93,7 @@ class AbstractOperator(LoggingMixin, DAGNode):
  # For derived classes to define which fields will get jinjaified.
  template_fields: Collection[str]
  # Defines which files extensions to look for in the templated fields.
- template_ext: Collection[str]
+ template_ext: Sequence[str]
 
  owner: str
  task_id: str

diff --git a/airflow/models/baseoperator.py b/airflow/models/baseoperator.py
@@ -596,8 +596,8 @@ class derived from this one results in the creation of a task object,
  """
 
  # Implementing Operator.
- template_fields: Collection[str] = ()
- template_ext: Collection[str] = ()
+ template_fields: Sequence[str] = ()
+ template_ext: Sequence[str] = ()
 
  template_fields_renderers: Dict[str, str] = {}
 

diff --git a/airflow/models/mappedoperator.py b/airflow/models/mappedoperator.py
@@ -250,7 +250,7 @@ class MappedOperator(AbstractOperator):
  params: Optional[dict]
  deps: FrozenSet[BaseTIDep]
  operator_extra_links: Collection["BaseOperatorLink"]
- template_ext: Collection[str]
+ template_ext: Sequence[str]
  template_fields: Collection[str]
  template_fields_renderers: Dict[str, str]
  ui_color: str

diff --git a/airflow/operators/bash.py b/airflow/operators/bash.py
@@ -21,7 +21,7 @@
 from airflow.compat.functools import cached_property
 from airflow.exceptions import AirflowException, AirflowSkipException
 from airflow.hooks.subprocess import SubprocessHook
-from airflow.models import BaseOperator
+from airflow.models.baseoperator import BaseOperator
 from airflow.utils.context import Context
 from airflow.utils.operator_helpers import context_to_airflow_vars
 

diff --git a/airflow/operators/dummy.py b/airflow/operators/dummy.py
@@ -15,7 +15,7 @@
 # KIND, either express or implied. See the License for the
 # specific language governing permissions and limitations
 # under the License.
-from airflow.models import BaseOperator
+from airflow.models.baseoperator import BaseOperator
 from airflow.utils.context import Context
 
 

diff --git a/airflow/operators/python.py b/airflow/operators/python.py
@@ -29,7 +29,7 @@
 import dill
 
 from airflow.exceptions import AirflowException
-from airflow.models import BaseOperator
+from airflow.models.baseoperator import BaseOperator
 from airflow.models.skipmixin import SkipMixin
 from airflow.models.taskinstance import _CURRENT_CONTEXT
 from airflow.utils.context import Context, context_copy_partial, context_merge

diff --git a/airflow/providers/amazon/aws/transfers/redshift_to_s3.py b/airflow/providers/amazon/aws/transfers/redshift_to_s3.py
@@ -139,8 +139,6 @@ def _build_unload_query(
  def execute(self, context: 'Context') -> None:
  redshift_hook = RedshiftSQLHook(redshift_conn_id=self.redshift_conn_id)
  conn = S3Hook.get_connection(conn_id=self.aws_conn_id)
-
- credentials_block = None
  if conn.extra_dejson.get('role_arn', False):
  credentials_block = f"aws_iam_role={conn.extra_dejson['role_arn']}"
  else:

diff --git a/airflow/providers/amazon/aws/transfers/s3_to_redshift.py b/airflow/providers/amazon/aws/transfers/s3_to_redshift.py
@@ -127,7 +127,6 @@ def execute(self, context: 'Context') -> None:
  redshift_hook = RedshiftSQLHook(redshift_conn_id=self.redshift_conn_id)
  conn = S3Hook.get_connection(conn_id=self.aws_conn_id)
 
- credentials_block = None
  if conn.extra_dejson.get('role_arn', False):
  credentials_block = f"aws_iam_role={conn.extra_dejson['role_arn']}"
  else:

diff --git a/airflow/providers/asana/hooks/asana.py b/airflow/providers/asana/hooks/asana.py
@@ -20,8 +20,8 @@
 import sys
 from typing import Any, Dict, Optional
 
-from asana import Client
-from asana.error import NotFoundError
+from asana import Client # type: ignore[attr-defined]
+from asana.error import NotFoundError # type: ignore[attr-defined]
 
 if sys.version_info >= (3, 8):
  from functools import cached_property

diff --git a/airflow/providers/cloudant/hooks/cloudant.py b/airflow/providers/cloudant/hooks/cloudant.py
@@ -18,7 +18,7 @@
 """Hook for Cloudant"""
 from typing import Any, Dict
 
-from cloudant import cloudant
+from cloudant import cloudant # type: ignore[attr-defined]
 
 from airflow.exceptions import AirflowException
 from airflow.hooks.base import BaseHook

diff --git a/airflow/providers/cncf/kubernetes/operators/kubernetes_pod.py b/airflow/providers/cncf/kubernetes/operators/kubernetes_pod.py
@@ -42,7 +42,7 @@
  convert_volume,
  convert_volume_mount,
 )
-from airflow.providers.cncf.kubernetes.utils import xcom_sidecar
+from airflow.providers.cncf.kubernetes.utils import xcom_sidecar # type: ignore[attr-defined]
 from airflow.providers.cncf.kubernetes.utils.pod_manager import PodLaunchFailedException, PodManager, PodPhase
 from airflow.settings import pod_mutation_hook
 from airflow.utils import yaml

diff --git a/airflow/providers/cncf/kubernetes/utils/__init__.py b/airflow/providers/cncf/kubernetes/utils/__init__.py
@@ -14,3 +14,4 @@
 # KIND, either express or implied. See the License for the
 # specific language governing permissions and limitations
 # under the License.
+__all__ = ['xcom_sidecar', 'pod_manager']
diff --git a/airflow/providers/databricks/hooks/databricks_sql.py b/airflow/providers/databricks/hooks/databricks_sql.py
@@ -20,8 +20,8 @@
 from copy import copy
 from typing import Any, Dict, List, Optional, Union
 
-from databricks import sql
-from databricks.sql.client import Connection
+from databricks import sql # type: ignore[attr-defined]
+from databricks.sql.client import Connection # type: ignore[attr-defined]
 
 from airflow.exceptions import AirflowException
 from airflow.hooks.dbapi import DbApiHook

diff --git a/airflow/providers/datadog/hooks/datadog.py b/airflow/providers/datadog/hooks/datadog.py
@@ -19,7 +19,7 @@
 import time
 from typing import Any, Dict, List, Optional, Union
 
-from datadog import api, initialize
+from datadog import api, initialize # type: ignore[attr-defined]
 
 from airflow.exceptions import AirflowException
 from airflow.hooks.base import BaseHook

diff --git a/airflow/providers/dbt/cloud/example_dags/example_dbt_cloud.py b/airflow/providers/dbt/cloud/example_dags/example_dbt_cloud.py
@@ -17,7 +17,7 @@
 
 from datetime import datetime
 
-from airflow.models import DAG
+from airflow.models import DAG, BaseOperator
 from airflow.operators.dummy import DummyOperator
 from airflow.providers.dbt.cloud.operators.dbt import (
  DbtCloudGetJobRunArtifactOperator,
@@ -46,7 +46,7 @@
  # [END howto_operator_dbt_cloud_run_job]
 
  # [START howto_operator_dbt_cloud_get_artifact]
- get_run_results_artifact = DbtCloudGetJobRunArtifactOperator(
+ get_run_results_artifact: BaseOperator = DbtCloudGetJobRunArtifactOperator(
  task_id="get_run_results_artifact", run_id=trigger_job_run1.output, path="run_results.json"
  )
  # [END howto_operator_dbt_cloud_get_artifact]
@@ -61,7 +61,7 @@
  # [END howto_operator_dbt_cloud_run_job_async]
 
  # [START howto_operator_dbt_cloud_run_job_sensor]
- job_run_sensor = DbtCloudJobRunSensor(
+ job_run_sensor: BaseOperator = DbtCloudJobRunSensor(
  task_id="job_run_sensor", run_id=trigger_job_run2.output, timeout=20
  )
  # [END howto_operator_dbt_cloud_run_job_sensor]

diff --git a/airflow/providers/docker/hooks/docker.py b/airflow/providers/docker/hooks/docker.py
@@ -17,9 +17,9 @@
 # under the License.
 from typing import Any, Dict, Optional
 
-from docker import APIClient
-from docker.constants import DEFAULT_TIMEOUT_SECONDS
-from docker.errors import APIError
+from docker import APIClient # type: ignore[attr-defined]
+from docker.constants import DEFAULT_TIMEOUT_SECONDS # type: ignore[attr-defined]
+from docker.errors import APIError # type: ignore[attr-defined]
 
 from airflow.exceptions import AirflowException
 from airflow.hooks.base import BaseHook

diff --git a/airflow/providers/docker/operators/docker.py b/airflow/providers/docker/operators/docker.py
@@ -23,10 +23,10 @@
 from tempfile import TemporaryDirectory
 from typing import TYPE_CHECKING, Dict, Iterable, List, Optional, Sequence, Union
 
-from docker import APIClient, tls
-from docker.constants import DEFAULT_TIMEOUT_SECONDS
-from docker.errors import APIError
-from docker.types import Mount
+from docker import APIClient, tls # type: ignore[attr-defined]
+from docker.constants import DEFAULT_TIMEOUT_SECONDS # type: ignore[attr-defined]
+from docker.errors import APIError # type: ignore[attr-defined]
+from docker.types import Mount # type: ignore[attr-defined]
 
 from airflow.exceptions import AirflowException
 from airflow.models import BaseOperator

diff --git a/airflow/providers/docker/operators/docker_swarm.py b/airflow/providers/docker/operators/docker_swarm.py
@@ -144,7 +144,8 @@ def _run_service(self) -> None:
  labels={'name': f'airflow__{self.dag_id}__{self.task_id}'},
  mode=self.mode,
  )
-
+ if self.service is None:
+ raise Exception("Service should be set here")
  self.log.info('Service started: %s', str(self.service))
 
  # wait for the service to start the task
@@ -171,6 +172,8 @@ def _run_service(self) -> None:
  def _service_status(self) -> Optional[str]:
  if not self.cli:
  raise Exception("The 'cli' should be initialized before!")
+ if not self.service:
+ raise Exception("The 'service' should be initialized before!")
  return self.cli.tasks(filters={'service': self.service['ID']})[0]['Status']['State']
 
  def _has_service_terminated(self) -> bool:
@@ -207,6 +210,6 @@ def _stream_logs_to_output(self) -> None:
  self.log.info(line)
 
  def on_kill(self) -> None:
- if self.cli is not None:
+ if self.cli is not None and self.service is not None:
  self.log.info('Removing docker service: %s', self.service['ID'])
  self.cli.remove_service(self.service['ID'])
diff --git a/airflow/providers/elasticsearch/log/es_task_handler.py b/airflow/providers/elasticsearch/log/es_task_handler.py
@@ -86,7 +86,7 @@ def __init__(
  super().__init__(base_log_folder, filename_template)
  self.closed = False
 
- self.client = elasticsearch.Elasticsearch([host], **es_kwargs)
+ self.client = elasticsearch.Elasticsearch([host], **es_kwargs) # type: ignore[attr-defined]
 
  self.log_id_template = log_id_template
  self.frontend = frontend

diff --git a/airflow/providers/github/example_dags/example_github.py b/airflow/providers/github/example_dags/example_github.py
@@ -56,7 +56,7 @@ def tag_checker(repo: Any, tag_name: str) -> Optional[bool]:
  all_tags = [x.name for x in repo.get_tags()]
  result = tag_name in all_tags
 
- except GithubException as github_error:
+ except GithubException as github_error: # type: ignore[misc]
  raise AirflowException(f"Failed to execute GithubSensor, error: {str(github_error)}")
  except Exception as e:
  raise AirflowException(f"Github operator error: {str(e)}")

diff --git a/airflow/providers/github/sensors/github.py b/airflow/providers/github/sensors/github.py
@@ -138,7 +138,7 @@ def tag_checker(self, repo: Any) -> Optional[bool]:
  all_tags = [x.name for x in repo.get_tags()]
  result = self.tag_name in all_tags
 
- except GithubException as github_error:
+ except GithubException as github_error: # type: ignore[misc]
  raise AirflowException(f"Failed to execute GithubSensor, error: {str(github_error)}")
  except Exception as e:
  raise AirflowException(f"Github operator error: {str(e)}")

diff --git a/airflow/providers/google/cloud/example_dags/example_vision.py b/airflow/providers/google/cloud/example_dags/example_vision.py
@@ -57,6 +57,8 @@
 )
 
 # [START howto_operator_vision_retry_import]
+
+
 from google.api_core.retry import Retry # isort:skip
 
 # [END howto_operator_vision_retry_import]