Enable MyPy in strict mode

.circleci/config.yml

@@ -8,6 +8,7 @@ workflows:
   test:
     jobs:
       - static-checks
+      - mypy
       - test:
           name: test-python<< matrix.python_version >>
           matrix:
@@ -49,6 +50,31 @@ jobs:
             - ~/.cache/pre-commit
             - ~/.pyenv/versions/
 
+  mypy:
+    description: "Mypy"
+    executor:
+      name: docker-executor
+      python_version: "3.9"
+    steps:
+      - checkout
+      - restore_cache:
+          keys:
+            - mypy-{{ .Branch }}-{{ checksum "setup.cfg" }}-{{ checksum "/home/circleci/.pyenv/version" }}
+            - mypy-main-{{ checksum "setup.cfg" }}-{{ checksum "/home/circleci/.pyenv/version" }}
+      - run:
+          name: Install Dependencies
+          command: pip install -U -e .[mypy]
+      - run:
+          name: Run Mypy
+          command: |
+            mypy --version
+            mypy astronomer/
+      - save_cache:
+          paths:
+            - ~/.cache/pip
+            - ~/.pyenv/versions/
+          key: mypy-{{ .Branch }}-{{ checksum "setup.cfg" }}-{{ checksum "/home/circleci/.pyenv/version" }}
+
   test:
     parameters:
       python_version:

.gitignore

@@ -42,3 +42,6 @@ docs/*/_api/
 test-report/
 .coverage
 coverage.xml
+
+# Mypy Cache
+.mypy_cache

.pre-commit-config.yaml

@@ -27,6 +27,7 @@ repos:
       - id: mixed-line-ending
       - id: trailing-whitespace
 
+
   - repo: https://github.com/codespell-project/codespell
     rev: v2.1.0
     hooks:

Makefile

@@ -49,7 +49,8 @@ run-mypy: ## Run MyPy in Container
 	docker build -f dev/Dockerfile . -t astronomer-providers-dev
 	docker run -v `pwd`:/usr/local/airflow/astronomer_providers -v `pwd`/dev/.cache:/home/astro/.cache \
 	 	-w /usr/local/airflow/astronomer_providers \
-		--rm -it astronomer-providers-dev -- mypy --install-types $(RUN_ARGS)
+		--rm -it astronomer-providers-dev \
+		-- mypy --install-types --cache-dir /home/astro/.cache/.mypy_cache $(RUN_ARGS)
 
 help: ## Prints this message
 	@grep -E '^[a-zA-Z_-]+:.*?## .*$$' $(MAKEFILE_LIST) | sort | awk 'BEGIN {FS = ":.*?## "}; {printf "\033[36m%-30s\033[0m %s\n", $$1, $$2}'

astronomer/providers/amazon/aws/hooks/s3.py

@@ -205,14 +205,14 @@ async def is_keys_unchanged(
 
             if current_num_objects >= min_objects:
                 success_message = (
-                    "SUCCESS: \nSensor found %s objects at %s.\n"
-                    "Waited at least %s seconds, with no new objects uploaded.",
-                    current_num_objects,
-                    path,
-                    inactivity_period,
+                    "SUCCESS: Sensor found %s objects at %s. "
+                    "Waited at least %s seconds, with no new objects uploaded."
                 )
-                self.log.info(success_message)
-                return {"status": "success", "message": success_message}
+                self.log.info(success_message, current_num_objects, path, inactivity_period)
+                return {
+                    "status": "success",
+                    "message": success_message % (current_num_objects, path, inactivity_period),
+                }
 
             self.log.error("FAILURE: Inactivity Period passed, not enough objects found in %s", path)
             return {

astronomer/providers/amazon/aws/operators/redshift_cluster.py

@@ -52,7 +52,7 @@ def execute(self, context: "Context") -> None:
                 "Unable to resume cluster since cluster is currently in status: %s", cluster_state
             )
 
-    def execute_complete(self, context: Dict[Any, Any], event: Any = None) -> None:
+    def execute_complete(self, context: Dict[str, Any], event: Any = None) -> None:
         """
         Callback for when the trigger fires - returns immediately.
         Relies on trigger to throw an exception, otherwise it assumes execution was
@@ -108,7 +108,7 @@ def execute(self, context: "Context") -> None:
                 "Unable to pause cluster since cluster is currently in status: %s", cluster_state
             )
 
-    def execute_complete(self, context: Dict[Any, Any], event: Any = None) -> None:
+    def execute_complete(self, context: Dict[str, Any], event: Any = None) -> None:
         """
         Callback for when the trigger fires - returns immediately.
         Relies on trigger to throw an exception, otherwise it assumes execution was

astronomer/providers/amazon/aws/operators/redshift_sql.py

@@ -1,11 +1,14 @@
-from typing import Any, Dict, cast
+from typing import TYPE_CHECKING, Any, Dict, cast
 
 from airflow.exceptions import AirflowException
 from airflow.providers.amazon.aws.operators.redshift_sql import RedshiftSQLOperator
 
 from astronomer.providers.amazon.aws.hooks.redshift_data import RedshiftDataHook
 from astronomer.providers.amazon.aws.triggers.redshift_sql import RedshiftSQLTrigger
 
+if TYPE_CHECKING:
+    from airflow.utils.context import Context
+
 
 class RedshiftSQLOperatorAsync(RedshiftSQLOperator):
     """
@@ -21,7 +24,7 @@ def __init__(
         self.poll_interval = poll_interval
         super().__init__(**kwargs)
 
-    def execute(self, context: Dict[str, Any]) -> None:
+    def execute(self, context: "Context") -> None:
         redshift_data_hook = RedshiftDataHook(aws_conn_id=self.redshift_conn_id)
         query_ids, response = redshift_data_hook.execute_query(sql=cast(str, self.sql), params=self.params)
         if response.get("status") == "error":

astronomer/providers/amazon/aws/sensors/redshift_cluster.py

@@ -1,5 +1,5 @@
 import logging
-from typing import Any, Dict
+from typing import TYPE_CHECKING, Any, Dict, Optional
 
 from airflow.exceptions import AirflowException
 from airflow.providers.amazon.aws.sensors.redshift_cluster import RedshiftClusterSensor
@@ -8,6 +8,9 @@
     RedshiftClusterSensorTrigger,
 )
 
+if TYPE_CHECKING:
+    from airflow.utils.context import Context
+
 log = logging.getLogger(__name__)
 
 
@@ -28,7 +31,7 @@ def __init__(
         self.poll_interval = poll_interval
         super().__init__(**kwargs)
 
-    def execute(self, context: Dict[Any, Any]) -> None:
+    def execute(self, context: Dict[str, Any]) -> None:
         self.defer(
             timeout=self.execution_timeout,
             trigger=RedshiftClusterSensorTrigger(
@@ -41,7 +44,7 @@ def execute(self, context: Dict[Any, Any]) -> None:
             method_name="execute_complete",
         )
 
-    def execute_complete(self, context: Dict[Any, Any], event: Any = None) -> None:
+    def execute_complete(self, context: "Context", event: Optional[Dict[Any, Any]] = None) -> None:
         """
         Callback for when the trigger fires - returns immediately.
         Relies on trigger to throw an exception, otherwise it assumes execution was

astronomer/providers/amazon/aws/sensors/s3.py

@@ -75,7 +75,7 @@ def _resolve_bucket_and_key(self) -> None:
             if parsed_url.scheme != "" or parsed_url.netloc != "":
                 raise AirflowException("If bucket_name provided, bucket_key must be relative path, not URI.")
 
-    def execute(self, context: Dict[Any, Any]) -> None:
+    def execute(self, context: Dict[str, Any]) -> None:
         self._resolve_bucket_and_key()
         self.defer(
             timeout=self.execution_timeout,
@@ -89,7 +89,7 @@ def execute(self, context: Dict[Any, Any]) -> None:
             method_name="execute_complete",
         )
 
-    def execute_complete(self, context: Dict[Any, Any], event: Any = None) -> None:
+    def execute_complete(self, context: Dict[str, Any], event: Any = None) -> None:
         if event["status"] == "error":
             raise AirflowException(event["message"])
         return None
@@ -137,7 +137,7 @@ def __init__(
         super().__init__(**kwargs)
         self.check_fn_user = check_fn
 
-    def execute(self, context: Dict[Any, Any]) -> None:
+    def execute(self, context: Dict[str, Any]) -> None:
         self._resolve_bucket_and_key()
         self.defer(
             timeout=self.execution_timeout,
@@ -152,7 +152,7 @@ def execute(self, context: Dict[Any, Any]) -> None:
             method_name="execute_complete",
         )
 
-    def execute_complete(self, context: Dict[Any, Any], event: Any = None) -> None:
+    def execute_complete(self, context: Dict[str, Any], event: Any = None) -> None:
         if event["status"] == "error":
             raise AirflowException(event["message"])
         return None
@@ -219,7 +219,7 @@ def __init__(
         self.verify = verify
         self.last_activity_time: Optional[datetime] = None
 
-    def execute(self, context: Dict[Any, Any]) -> None:
+    def execute(self, context: Dict[str, Any]) -> None:
         self.defer(
             timeout=self.execution_timeout,
             trigger=S3KeysUnchangedTrigger(
@@ -237,7 +237,7 @@ def execute(self, context: Dict[Any, Any]) -> None:
             method_name="execute_complete",
         )
 
-    def execute_complete(self, context: Dict[Any, Any], event: Any = None) -> None:
+    def execute_complete(self, context: Dict[str, Any], event: Any = None) -> None:
         if event["status"] == "error":
             raise AirflowException(event["message"])
         return None

astronomer/providers/cncf/kubernetes/hooks/kubernetes_async.py

@@ -1,11 +1,11 @@
 import aiofiles
-from airflow import AirflowException
+from airflow.exceptions import AirflowException
 from airflow.providers.cncf.kubernetes.hooks.kubernetes import KubernetesHook
 from kubernetes_asyncio import client, config
 
 
 class KubernetesHookAsync(KubernetesHook):
-    async def _load_config(self):
+    async def _load_config(self) -> client.ApiClient:
         """
         cluster_context: Optional[str] = None,
         config_file: Optional[str] = None,
@@ -50,7 +50,7 @@ async def _load_config(self):
             return client.ApiClient()
 
         if kubeconfig is not None:
-            async with aiofiles.tempfile.NamedTemporaryFile() as temp_config:
+            async with aiofiles.tempfile.NamedTemporaryFile() as temp_config:  # type: ignore[attr-defined]
                 self.log.debug("loading kube_config from: connection kube_config")
                 temp_config.write(kubeconfig.encode())
                 temp_config.flush()

astronomer/providers/cncf/kubernetes/operators/kubernetes_pod.py

@@ -1,7 +1,10 @@
-from airflow import AirflowException
+from typing import Any, Dict
+
+from airflow.exceptions import AirflowException
 from airflow.providers.cncf.kubernetes.operators.kubernetes_pod import (
     KubernetesPodOperator,
 )
+from airflow.utils.context import Context
 
 from astronomer.providers.cncf.kubernetes.triggers.wait_container import (
     PodLaunchTimeoutException,
@@ -20,12 +23,14 @@ class KubernetesPodOperatorAsync(KubernetesPodOperator):
     :param poll_interval: interval in seconds to sleep between checking pod status
     """
 
-    def __init__(self, *, poll_interval: int = 5, **kwargs):
+    def __init__(self, *, poll_interval: int = 5, **kwargs: Any):
+        self.pod = None
+        self.pod_request_obj = None
         self.poll_interval = poll_interval
         super().__init__(**kwargs)
 
     @staticmethod
-    def raise_for_trigger_status(event):
+    def raise_for_trigger_status(event: Dict[str, Any]) -> None:
         if event["status"] == "error":
             error_type = event["error_type"]
             description = event["description"]
@@ -34,7 +39,7 @@ def raise_for_trigger_status(event):
             else:
                 raise AirflowException(description)
 
-    def execute(self, context):
+    def execute(self, context: Context) -> None:
         self.pod_request_obj = self.build_pod_request_obj(context)
         self.pod = self.get_or_create_pod(self.pod_request_obj, context)
         self.defer(
@@ -54,7 +59,7 @@ def execute(self, context):
             method_name=self.execute_complete.__name__,
         )
 
-    def execute_complete(self, context, event=None):
+    def execute_complete(self, context: Context, event: Dict[str, Any]) -> Any:
         remote_pod = None
         try:
             self.pod_request_obj = self.build_pod_request_obj(context)

astronomer/providers/cncf/kubernetes/triggers/wait_container.py

@@ -1,7 +1,7 @@
 import asyncio
 import traceback
 from datetime import timedelta
-from typing import Any, Dict, Optional, Tuple
+from typing import Any, AsyncIterator, Dict, Optional, Tuple
 
 from airflow.exceptions import AirflowException
 from airflow.providers.cncf.kubernetes.utils.pod_manager import (
@@ -27,32 +27,27 @@ class WaitContainerTrigger(BaseTrigger):
     Next, waits for ``container_name`` to reach a terminal state.
 
     :param kubernetes_conn_id: Airflow connection ID to use
-    :type kubernetes_conn_id: str
     :param hook_params: kwargs for hook
-    :type hook_params: dict
     :param container_name: container to wait for
-    :type container_name: str
     :param pod_name: name of pod to monitor
-    :type pod_name: str
     :param pod_namespace: pod namespace
-    :type pod_namespace: str
     :param pending_phase_timeout: max time in seconds to wait for pod to leave pending phase
-    :type pending_phase_timeout: float
     :param poll_interval: number of seconds between reading pod state
-    :type poll_interval: float
 
     """
 
     def __init__(
         self,
+        *,
+        container_name: str,
+        pod_name: str,
+        pod_namespace: str,
         kubernetes_conn_id: Optional[str] = None,
-        hook_params: Optional[dict] = None,
-        container_name: Optional[str] = None,
-        pod_name: Optional[str] = None,
-        pod_namespace: Optional[str] = None,
+        hook_params: Optional[Dict[str, Any]] = None,
         pending_phase_timeout: float = 120,
         poll_interval: float = 5,
     ):
+        super().__init__()
         self.kubernetes_conn_id = kubernetes_conn_id
         self.hook_params = hook_params
         self.container_name = container_name
@@ -78,7 +73,7 @@ def serialize(self) -> Tuple[str, Dict[str, Any]]:
     async def get_hook(self) -> KubernetesHookAsync:
         return KubernetesHookAsync(conn_id=self.kubernetes_conn_id, **(self.hook_params or {}))
 
-    async def wait_for_pod_start(self, v1_api: CoreV1Api):
+    async def wait_for_pod_start(self, v1_api: CoreV1Api) -> Any:
         """
         Loops until pod phase leaves ``PENDING``
         If timeout is reached, throws error.
@@ -92,15 +87,15 @@ async def wait_for_pod_start(self, v1_api: CoreV1Api):
             await asyncio.sleep(self.poll_interval)
         raise PodLaunchTimeoutException("Pod did not leave 'Pending' phase within specified timeout")
 
-    async def wait_for_container_completion(self, v1_api: CoreV1Api):
+    async def wait_for_container_completion(self, v1_api: CoreV1Api) -> None:
         """Waits until container ``self.container_name`` is no longer in running state."""
         while True:
             pod = await v1_api.read_namespaced_pod(self.pod_name, self.pod_namespace)
             if not container_is_running(pod=pod, container_name=self.container_name):
                 break
             await asyncio.sleep(self.poll_interval)
 
-    async def run(self):
+    async def run(self) -> AsyncIterator["TriggerEvent"]:  # type: ignore[override]
         self.log.debug("Checking pod %r in namespace %r.", self.pod_name, self.pod_namespace)
         try:
             hook = await self.get_hook()
@@ -120,7 +115,7 @@ async def run(self):
                 }
             )
 
-    def _format_exception_description(self, exc: Exception):
+    def _format_exception_description(self, exc: Exception) -> Any:
         if isinstance(exc, PodLaunchTimeoutException):
             return exc.args[0]