Upgrade mypy to 0.991 (#28926)

ashb · uranusjr · web-flow · commit 6ae0a80cbaf1 · 2023-02-11T02:25:39.000+08:00
Co-authored-by: Tzu-ping Chung &lt;uranusjr@gmail.com&gt;
diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml
@@ -160,7 +160,7 @@ repos:
         entry: ./scripts/ci/pre_commit/pre_commit_update_common_sql_api_stubs.py
         language: python
         files: ^scripts/ci/pre_commit/pre_commit_update_common_sql_api\.py|^airflow/providers/common/sql/.*\.pyi?$
-        additional_dependencies: ['rich>=12.4.4', 'mypy==0.971', 'black==22.12.0', 'jinja2']
+        additional_dependencies: ['rich>=12.4.4', 'mypy==0.991', 'black==22.12.0', 'jinja2']
         pass_filenames: false
         require_serial: true
       - id: update-black-version
diff --git a/airflow/dag_processing/manager.py b/airflow/dag_processing/manager.py
@@ -141,8 +141,7 @@ def __init__(
 
     def start(self) -> None:
         """Launch DagFileProcessorManager processor and start DAG parsing loop in manager."""
-        mp_start_method = self._get_multiprocessing_start_method()
-        context = multiprocessing.get_context(mp_start_method)
+        context = self._get_multiprocessing_context()
         self._last_parsing_stat_received_at = time.monotonic()
 
         self._parent_signal_conn, child_signal_conn = context.Pipe()
diff --git a/airflow/dag_processing/processor.py b/airflow/dag_processing/processor.py
@@ -187,8 +187,7 @@ def _handle_dag_file_processing():
 
     def start(self) -> None:
         """Launch the process and start processing the DAG."""
-        start_method = self._get_multiprocessing_start_method()
-        context = multiprocessing.get_context(start_method)
+        context = self._get_multiprocessing_context()
 
         _parent_channel, _child_channel = context.Pipe(duplex=False)
         process = context.Process(
diff --git a/airflow/decorators/base.py b/airflow/decorators/base.py
@@ -489,7 +489,7 @@ def _get_unmap_kwargs(self, mapped_kwargs: Mapping[str, Any], *, strict: bool) -
         return super()._get_unmap_kwargs(kwargs, strict=False)
 
 
-class Task(Generic[FParams, FReturn]):
+class Task(Protocol, Generic[FParams, FReturn]):
     """Declaration of a @task-decorated callable for type-checking.
 
     An instance of this type inherits the call signature of the decorated
diff --git a/airflow/example_dags/plugins/workday.py b/airflow/example_dags/plugins/workday.py
@@ -78,6 +78,9 @@ def next_dagrun_info(
 
     # [END howto_timetable_next_dagrun_info]
 
+    def validate(self):
+        pass
+
 
 class WorkdayTimetablePlugin(AirflowPlugin):
     name = "workday_timetable_plugin"
diff --git a/airflow/executors/kubernetes_executor.py b/airflow/executors/kubernetes_executor.py
@@ -69,7 +69,7 @@
 class ResourceVersion:
     """Singleton for tracking resourceVersion from Kubernetes."""
 
-    _instance = None
+    _instance: ResourceVersion | None = None
     resource_version: dict[str, str] = {}
 
     def __new__(cls):
diff --git a/airflow/hooks/base.py b/airflow/hooks/base.py
@@ -90,11 +90,11 @@ def get_conn(self) -> Any:
 
     @classmethod
     def get_connection_form_widgets(cls) -> dict[str, Any]:
-        ...
+        return {}
 
     @classmethod
     def get_ui_field_behaviour(cls) -> dict[str, Any]:
-        ...
+        return {}
 
 
 class DiscoverableHook(Protocol):
diff --git a/airflow/listeners/listener.py b/airflow/listeners/listener.py
@@ -30,7 +30,7 @@
 log = logging.getLogger(__name__)
 
 
-_listener_manager = None
+_listener_manager: ListenerManager | None = None
 
 
 class ListenerManager:
diff --git a/airflow/models/baseoperator.py b/airflow/models/baseoperator.py
@@ -30,7 +30,7 @@
 from abc import ABCMeta, abstractmethod
 from datetime import datetime, timedelta
 from inspect import signature
-from types import FunctionType
+from types import ClassMethodDescriptorType, FunctionType
 from typing import (
     TYPE_CHECKING,
     Any,
@@ -169,7 +169,7 @@ def get_merged_defaults(
 class _PartialDescriptor:
     """A descriptor that guards against ``.partial`` being called on Task objects."""
 
-    class_method = None
+    class_method: ClassMethodDescriptorType | None = None
 
     def __get__(
         self, obj: BaseOperator, cls: type[BaseOperator] | None = None
diff --git a/airflow/providers/google/cloud/hooks/cloud_sql.py b/airflow/providers/google/cloud/hooks/cloud_sql.py
@@ -691,8 +691,6 @@ class CloudSQLDatabaseHook(BaseHook):
     conn_type = "gcpcloudsqldb"
     hook_name = "Google Cloud SQL Database"
 
-    _conn = None
-
     def __init__(
         self,
         gcp_cloudsql_conn_id: str = "google_cloud_sql_default",
diff --git a/airflow/providers/google/cloud/hooks/functions.py b/airflow/providers/google/cloud/hooks/functions.py
@@ -40,7 +40,7 @@ class CloudFunctionsHook(GoogleBaseHook):
     keyword arguments rather than positional.
     """
 
-    _conn = None
+    _conn: build | None = None
 
     def __init__(
         self,
diff --git a/airflow/providers/google/cloud/hooks/gcs.py b/airflow/providers/google/cloud/hooks/gcs.py
@@ -51,8 +51,18 @@
 from airflow.utils import timezone
 from airflow.version import version
 
+try:
+    # Airflow 2.3 doesn't have this yet
+    from airflow.typing_compat import ParamSpec
+except ImportError:
+    try:
+        from typing import ParamSpec  # type: ignore[no-redef,attr-defined]
+    except ImportError:
+        from typing_extensions import ParamSpec
+
 RT = TypeVar("RT")
 T = TypeVar("T", bound=Callable)
+FParams = ParamSpec("FParams")
 
 # GCSHook has a method named 'list' (to junior devs: please don't do this), so
 # we need to create an alias to prevent Mypy being confused.
@@ -76,7 +86,7 @@ def _fallback_object_url_to_object_name_and_bucket_name(
     :return: Decorator
     """
 
-    def _wrapper(func: T):
+    def _wrapper(func: Callable[FParams, RT]) -> Callable[FParams, RT]:
         @functools.wraps(func)
         def _inner_wrapper(self: GCSHook, *args, **kwargs) -> RT:
             if args:
@@ -117,11 +127,11 @@ def _inner_wrapper(self: GCSHook, *args, **kwargs) -> RT:
                     f"'{bucket_name_keyword_arg_name}'"
                 )
 
-            return func(self, *args, **kwargs)
+            return func(self, *args, **kwargs)  # type: ignore
 
-        return cast(T, _inner_wrapper)
+        return cast(Callable[FParams, RT], _inner_wrapper)
 
-    return _wrapper
+    return cast(Callable[[T], T], _wrapper)
 
 
 # A fake bucket to use in functions decorated by _fallback_object_url_to_object_name_and_bucket_name.
diff --git a/airflow/providers/google/cloud/hooks/life_sciences.py b/airflow/providers/google/cloud/hooks/life_sciences.py
@@ -54,7 +54,7 @@ class LifeSciencesHook(GoogleBaseHook):
         account from the list granting this role to the originating account.
     """
 
-    _conn = None
+    _conn: build | None = None
 
     def __init__(
         self,
diff --git a/airflow/providers/google/cloud/utils/field_validator.py b/airflow/providers/google/cloud/utils/field_validator.py
@@ -206,7 +206,7 @@ def _sanity_checks(
         full_field_path: str,
         regexp: str,
         allow_empty: bool,
-        custom_validation: Callable,
+        custom_validation: Callable | None,
         value,
     ) -> None:
         if value is None and field_type != "union":
diff --git a/airflow/providers/google/firebase/hooks/firestore.py b/airflow/providers/google/firebase/hooks/firestore.py
@@ -53,7 +53,7 @@ class CloudFirestoreHook(GoogleBaseHook):
         account from the list granting this role to the originating account.
     """
 
-    _conn = None
+    _conn: build | None = None
 
     def __init__(
         self,
diff --git a/airflow/providers/google/marketing_platform/hooks/search_ads.py b/airflow/providers/google/marketing_platform/hooks/search_ads.py
@@ -28,7 +28,7 @@
 class GoogleSearchAdsHook(GoogleBaseHook):
     """Hook for Google Search Ads 360."""
 
-    _conn = None
+    _conn: build | None = None
 
     def __init__(
         self,
diff --git a/airflow/providers/microsoft/psrp/hooks/psrp.py b/airflow/providers/microsoft/psrp/hooks/psrp.py
@@ -75,8 +75,7 @@ class PsrpHook(BaseHook):
     or by setting this key as the extra fields of your connection.
     """
 
-    _conn = None
-    _configuration_name = None
+    _conn: RunspacePool | None = None
     _wsman_ref: WeakKeyDictionary[RunspacePool, WSMan] = WeakKeyDictionary()
 
     def __init__(
diff --git a/airflow/settings.py b/airflow/settings.py
@@ -411,7 +411,7 @@ def dispose_orm():
     global engine
     global Session
 
-    if Session:
+    if Session is not None:
         Session.remove()
         Session = None
     if engine:
diff --git a/airflow/stats.py b/airflow/stats.py
@@ -74,9 +74,10 @@ def timing(cls, stat: str, dt: float | datetime.timedelta, tags: dict[str, str]
     @classmethod
     def timer(cls, *args, **kwargs) -> TimerProtocol:
         """Timer metric that can be cancelled."""
+        raise NotImplementedError()
 
 
-class Timer:
+class Timer(TimerProtocol):
     """
     Timer that records duration, and optional sends to StatsD backend.
 
@@ -360,7 +361,7 @@ def timer(self, stat=None, *args, tags=None, **kwargs):
 
 
 class _Stats(type):
-    factory = None
+    factory: Callable[[], StatsLogger]
     instance: StatsLogger | None = None
 
     def __getattr__(cls, name):
@@ -374,7 +375,7 @@ def __getattr__(cls, name):
 
     def __init__(cls, *args, **kwargs):
         super().__init__(cls)
-        if cls.__class__.factory is None:
+        if not hasattr(cls.__class__, "factory"):
             is_datadog_enabled_defined = conf.has_option("metrics", "statsd_datadog_enabled")
             if is_datadog_enabled_defined and conf.getboolean("metrics", "statsd_datadog_enabled"):
                 cls.__class__.factory = cls.get_dogstatsd_logger
diff --git a/airflow/timetables/simple.py b/airflow/timetables/simple.py
@@ -55,6 +55,9 @@ def serialize(self) -> dict[str, Any]:
     def infer_manual_data_interval(self, *, run_after: DateTime) -> DataInterval:
         return DataInterval.exact(run_after)
 
+    def validate(self):
+        pass
+
 
 class NullTimetable(_TrivialTimetable):
     """Timetable that never schedules anything.
diff --git a/airflow/utils/mixins.py b/airflow/utils/mixins.py
@@ -25,6 +25,8 @@
 from airflow.utils.context import Context
 
 if typing.TYPE_CHECKING:
+    import multiprocessing.context
+
     from airflow.models.operator import Operator
 
 
@@ -44,6 +46,10 @@ def _get_multiprocessing_start_method(self) -> str:
             raise ValueError("Failed to determine start method")
         return method
 
+    def _get_multiprocessing_context(self) -> multiprocessing.context.DefaultContext:
+        mp_start_method = self._get_multiprocessing_start_method()
+        return multiprocessing.get_context(mp_start_method)  # type: ignore
+
 
 class ResolveMixin:
     """A runtime-resolved value."""
diff --git a/airflow/utils/session.py b/airflow/utils/session.py
@@ -28,7 +28,7 @@
 @contextlib.contextmanager
 def create_session() -> Generator[settings.SASession, None, None]:
     """Contextmanager that will create and teardown a session."""
-    if not settings.Session:
+    if not hasattr(settings, "Session") or settings.Session is None:
         raise RuntimeError("Session must be set before!")
     session = settings.Session()
     try:
diff --git a/airflow/utils/state.py b/airflow/utils/state.py
@@ -85,7 +85,7 @@ class State:
     FAILED = TaskInstanceState.FAILED
 
     # These are TaskState only
-    NONE = None
+    NONE: None = None
     REMOVED = TaskInstanceState.REMOVED
     SCHEDULED = TaskInstanceState.SCHEDULED
     QUEUED = TaskInstanceState.QUEUED
diff --git a/airflow/www/extensions/init_appbuilder.py b/airflow/www/extensions/init_appbuilder.py
@@ -14,6 +14,8 @@
 # KIND, either express or implied.  See the License for the
 # specific language governing permissions and limitations
 # under the License.
+
+# mypy: disable-error-code=var-annotated
 from __future__ import annotations
 
 import logging
diff --git a/airflow/www/fab_security/manager.py b/airflow/www/fab_security/manager.py
@@ -14,6 +14,8 @@
 # KIND, either express or implied.  See the License for the
 # specific language governing permissions and limitations
 # under the License.
+
+# mypy: disable-error-code=var-annotated
 from __future__ import annotations
 
 import base64
diff --git a/dev/breeze/src/airflow_breeze/utils/click_utils.py b/dev/breeze/src/airflow_breeze/utils/click_utils.py
@@ -19,4 +19,4 @@
 try:
     from rich_click import RichGroup as BreezeGroup
 except ImportError:
-    from click import Group as BreezeGroup  # type: ignore[misc] # noqa
+    from click import Group as BreezeGroup  # type: ignore[assignment] # noqa
diff --git a/scripts/ci/pre_commit/pre_commit_check_order_dockerfile_extras.py b/scripts/ci/pre_commit/pre_commit_check_order_dockerfile_extras.py
@@ -28,7 +28,7 @@
 
 from rich import print
 
-errors = []
+errors: list[str] = []
 
 MY_DIR_PATH = Path(__file__).parent.resolve()
 
diff --git a/scripts/ci/pre_commit/pre_commit_check_order_setup.py b/scripts/ci/pre_commit/pre_commit_check_order_setup.py
@@ -30,7 +30,7 @@
 
 from rich import print
 
-errors = []
+errors: list[str] = []
 
 MY_DIR_PATH = os.path.dirname(__file__)
 SOURCE_DIR_PATH = os.path.abspath(os.path.join(MY_DIR_PATH, os.pardir, os.pardir, os.pardir))
diff --git a/scripts/ci/pre_commit/pre_commit_check_system_tests_hidden_in_index.py b/scripts/ci/pre_commit/pre_commit_check_system_tests_hidden_in_index.py
@@ -19,6 +19,7 @@
 
 import sys
 from pathlib import Path
+from typing import Any
 
 from rich.console import Console
 
@@ -37,7 +38,7 @@
 PREFIX = "apache-airflow-providers-"
 
 
-errors = []
+errors: list[Any] = []
 
 
 def check_system_test_entry_hidden(provider_index: Path):
diff --git a/scripts/ci/pre_commit/pre_commit_update_common_sql_api_stubs.py b/scripts/ci/pre_commit/pre_commit_update_common_sql_api_stubs.py
diff --git a/scripts/in_container/run_provider_yaml_files_check.py b/scripts/in_container/run_provider_yaml_files_check.py
diff --git a/setup.cfg b/setup.cfg
diff --git a/setup.py b/setup.py
diff --git a/tests/listeners/dag_listener.py b/tests/listeners/dag_listener.py
diff --git a/tests/listeners/full_listener.py b/tests/listeners/full_listener.py
diff --git a/tests/listeners/lifecycle_listener.py b/tests/listeners/lifecycle_listener.py
diff --git a/tests/listeners/partial_listener.py b/tests/listeners/partial_listener.py
diff --git a/tests/test_utils/timetables.py b/tests/test_utils/timetables.py
diff --git a/tests/www/views/test_views_connection.py b/tests/www/views/test_views_connection.py
