Hook into Mypy to get rid of those cast() (#26023)

.pre-commit-config.yaml

@@ -239,6 +239,15 @@ repos:
         files: \.py$|\.pyi$
         # To keep consistent with the global isort skip config defined in setup.cfg
         exclude: ^airflow/_vendor/|^build/.*$|^venv/.*$|^\.tox/.*$
+        args:
+          # These -p options are duplicated to known_first_party in setup.cfg,
+          # Please keep these in sync for now. (See comments there for details.)
+          - -p=airflow
+          - -p=airflow_breeze
+          - -p=docker_tests
+          - -p=docs
+          - -p=kubernetes_tests
+          - -p=tests
   - repo: https://github.com/pycqa/pydocstyle
     rev: 6.1.1
     hooks:

airflow/mypy/plugin/outputs.py

@@ -0,0 +1,76 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+from __future__ import annotations
+
+from typing import Callable
+
+from mypy.plugin import AttributeContext, MethodContext, Plugin
+from mypy.types import AnyType, Type, TypeOfAny
+
+OUTPUT_PROPERTIES = {
+    "airflow.models.baseoperator.BaseOperator.output",
+    "airflow.models.mappedoperator.MappedOperator.output",
+}
+
+TASK_CALL_FUNCTIONS = {
+    "airflow.decorators.base.Task.__call__",
+}
+
+
+class OperatorOutputPlugin(Plugin):
+    """Plugin to convert XComArg to the runtime type.
+
+    This allows us to pass an *XComArg* to a downstream task, such as::
+
+        @task
+        def f(a: str) -> int:
+            return len(a)
+
+        f(op.output)  # "op" is an operator instance.
+        f(g())  # "g" is a taskflow task.
+
+    where the *a* argument of ``f`` should accept a *str* at runtime, but can be
+    provided with an *XComArg* in the DAG.
+
+    In the long run, it is probably a good idea to make *XComArg* a generic that
+    carries information about the task's return type, and build the entire XCom
+    mechanism into the type checker. But Python's type system is still limiting
+    in this regard now, and (using the above example) we yet to have a good way
+    to convert ``f``'s argument list from ``[str]`` to ``[XComArg[str] | str]``.
+    Perhaps *ParamSpec* will be extended enough one day to accommodate this.
+    """
+
+    @staticmethod
+    def _treat_as_any(context: AttributeContext | MethodContext) -> Type:
+        """Pretend *XComArg* is actually *typing.Any*."""
+        return AnyType(TypeOfAny.special_form, line=context.context.line, column=context.context.column)
+
+    def get_attribute_hook(self, fullname: str) -> Callable[[AttributeContext], Type] | None:
+        if fullname not in OUTPUT_PROPERTIES:
+            return None
+        return self._treat_as_any
+
+    def get_method_hook(self, fullname: str) -> Callable[[MethodContext], Type] | None:
+        if fullname not in TASK_CALL_FUNCTIONS:
+            return None
+        return self._treat_as_any
+
+
+def plugin(version: str):
+    return OperatorOutputPlugin

airflow/providers/amazon/aws/example_dags/example_ecs.py

@@ -16,7 +16,6 @@
 # under the License.
 
 from datetime import datetime
-from typing import cast
 
 from airflow import DAG
 from airflow.models.baseoperator import chain
@@ -100,7 +99,7 @@
     )
     # [END howto_operator_ecs_register_task_definition]
 
-    registered_task_definition = cast(str, register_task.output)
+    registered_task_definition = register_task.output
 
     # [START howto_sensor_ecs_task_definition_state]
     await_task_definition = EcsTaskDefinitionStateSensor(

airflow/providers/amazon/aws/example_dags/example_emr.py

@@ -17,7 +17,6 @@
 # under the License.
 import os
 from datetime import datetime
-from typing import cast
 
 from airflow import DAG
 from airflow.models.baseoperator import chain
@@ -80,7 +79,7 @@
     )
     # [END howto_operator_emr_create_job_flow]
 
-    job_flow_id = cast(str, job_flow_creator.output)
+    job_flow_id = job_flow_creator.output
 
     # [START howto_sensor_emr_job_flow]
     job_sensor = EmrJobFlowSensor(task_id='check_job_flow', job_flow_id=job_flow_id)

airflow/providers/google/cloud/example_dags/example_datafusion.py

@@ -20,7 +20,6 @@
 """
 import os
 from datetime import datetime
-from typing import cast
 
 from airflow import models
 from airflow.operators.bash import BashOperator
@@ -222,7 +221,7 @@
     start_pipeline_sensor = CloudDataFusionPipelineStateSensor(
         task_id="pipeline_state_sensor",
         pipeline_name=PIPELINE_NAME,
-        pipeline_id=cast(str, start_pipeline_async.output),
+        pipeline_id=start_pipeline_async.output,
         expected_statuses=["COMPLETED"],
         failure_statuses=["FAILED"],
         instance_name=INSTANCE_NAME,

airflow/providers/google/cloud/example_dags/example_looker.py

@@ -21,7 +21,6 @@
 """
 
 from datetime import datetime
-from typing import cast
 
 from airflow import models
 from airflow.providers.google.cloud.operators.looker import LookerStartPdtBuildOperator
@@ -44,7 +43,7 @@
     check_pdt_task_async_sensor = LookerCheckPdtBuildSensor(
         task_id='check_pdt_task_async_sensor',
         looker_conn_id='your_airflow_connection_for_looker',
-        materialization_id=cast(str, start_pdt_task_async.output),
+        materialization_id=start_pdt_task_async.output,
         poke_interval=10,
     )
     # [END cloud_looker_async_start_pdt_sensor]

airflow/providers/google/cloud/example_dags/example_vision.py

@@ -33,7 +33,6 @@
 
 import os
 from datetime import datetime
-from typing import cast
 
 from airflow import models
 from airflow.operators.bash import BashOperator
@@ -136,7 +135,7 @@
     )
     # [END howto_operator_vision_product_set_create]
 
-    product_set_create_output = cast(str, product_set_create.output)
+    product_set_create_output = product_set_create.output
 
     # [START howto_operator_vision_product_set_get]
     product_set_get = CloudVisionGetProductSetOperator(
@@ -173,7 +172,7 @@
     )
     # [END howto_operator_vision_product_create]
 
-    product_create_output = cast(str, product_create.output)
+    product_create_output = product_create.output
 
     # [START howto_operator_vision_product_get]
     product_get = CloudVisionGetProductOperator(

setup.cfg

@@ -190,7 +190,8 @@ no_implicit_optional = True
 warn_redundant_casts = True
 warn_unused_ignores = False
 plugins =
-  airflow.mypy.plugin.decorators
+  airflow.mypy.plugin.decorators,
+  airflow.mypy.plugin.outputs
 pretty = True
 show_error_codes = True
 
@@ -206,7 +207,10 @@ no_implicit_optional = False
 line_length=110
 combine_as_imports = true
 default_section = THIRDPARTY
-known_first_party=airflow,airflow_breeze,tests,docs
+# This is duplicated with arguments in .pre-commit-config.yaml because isort is
+# having some issues picking up these config files. Please keep these in sync
+# for now and track the isort issue: https://github.com/PyCQA/isort/issues/1889
+known_first_party = airflow,airflow_breeze,docker_tests,docs,kubernetes_tests,tests
 # Need to be consistent with the exclude config defined in pre-commit-config.yaml
-skip=build,.tox,venv
+skip = build,.tox,venv
 profile = black

tests/system/providers/airbyte/example_airbyte_trigger_job.py

@@ -20,7 +20,6 @@
 
 import os
 from datetime import datetime, timedelta
-from typing import cast
 
 from airflow import DAG
 from airflow.providers.airbyte.operators.airbyte import AirbyteTriggerSyncOperator
@@ -55,7 +54,7 @@
 
     airbyte_sensor = AirbyteJobSensor(
         task_id='airbyte_sensor_source_dest_example',
-        airbyte_job_id=cast(int, async_source_destination.output),
+        airbyte_job_id=async_source_destination.output,
     )
     # [END howto_operator_airbyte_asynchronous]
 

tests/system/providers/amazon/aws/example_athena.py

@@ -15,7 +15,6 @@
 # specific language governing permissions and limitations
 # under the License.
 from datetime import datetime
-from typing import cast
 
 import boto3
 
@@ -122,7 +121,7 @@ def read_results_from_s3(bucket_name, query_execution_id):
     # [START howto_sensor_athena]
     await_query = AthenaSensor(
         task_id='await_query',
-        query_execution_id=cast(str, read_table.output),
+        query_execution_id=read_table.output,
     )
     # [END howto_sensor_athena]
 

tests/system/providers/amazon/aws/example_batch.py

@@ -15,7 +15,6 @@
 # specific language governing permissions and limitations
 # under the License.
 from datetime import datetime
-from typing import cast
 
 import boto3
 
@@ -194,7 +193,7 @@ def delete_job_queue(job_queue_name):
     # [START howto_sensor_batch]
     wait_for_batch_job = BatchSensor(
         task_id='wait_for_batch_job',
-        job_id=cast(str, submit_batch_job.output),
+        job_id=submit_batch_job.output,
     )
     # [END howto_sensor_batch]
 

tests/system/providers/amazon/aws/example_emr_serverless.py

@@ -17,7 +17,6 @@
 
 
 from datetime import datetime
-from typing import cast
 
 from airflow.models.baseoperator import chain
 from airflow.models.dag import DAG
@@ -76,7 +75,7 @@
     )
     # [END howto_operator_emr_serverless_create_application]
 
-    emr_serverless_app_id = cast(str, emr_serverless_app.output)
+    emr_serverless_app_id = emr_serverless_app.output
 
     # [START howto_sensor_emr_serverless_application]
     wait_for_app_creation = EmrServerlessApplicationSensor(
@@ -97,7 +96,7 @@
 
     # [START howto_sensor_emr_serverless_job]
     wait_for_job = EmrServerlessJobSensor(
-        task_id='wait_for_job', application_id=emr_serverless_app_id, job_run_id=cast(str, start_job.output)
+        task_id='wait_for_job', application_id=emr_serverless_app_id, job_run_id=start_job.output
     )
     # [END howto_sensor_emr_serverless_job]
 

tests/system/providers/amazon/aws/example_glue.py

@@ -15,7 +15,7 @@
 # specific language governing permissions and limitations
 # under the License.
 from datetime import datetime
-from typing import List, Optional, Tuple, cast
+from typing import List, Optional, Tuple
 
 import boto3
 from botocore.client import BaseClient
@@ -66,7 +66,7 @@
 
 
 @task
-def get_role_name(arn):
+def get_role_name(arn: str) -> str:
     return arn.split('/')[-1]
 
 
@@ -162,7 +162,7 @@ def glue_cleanup(crawler_name: str, job_name: str, db_name: str) -> None:
         job_name=glue_job_name,
         script_location=f's3://{bucket_name}/etl_script.py',
         s3_bucket=bucket_name,
-        iam_role_name=cast(str, role_name),
+        iam_role_name=role_name,
         create_job_kwargs={'GlueVersion': '3.0', 'NumberOfWorkers': 2, 'WorkerType': 'G.1X'},
         # Waits by default, set False to test the Sensor below
         wait_for_completion=False,
@@ -174,7 +174,7 @@ def glue_cleanup(crawler_name: str, job_name: str, db_name: str) -> None:
         task_id='wait_for_job',
         job_name=glue_job_name,
         # Job ID extracted from previous Glue Job Operator task
-        run_id=cast(str, submit_glue_job.output),
+        run_id=submit_glue_job.output,
     )
     # [END howto_sensor_glue]
 
@@ -199,7 +199,7 @@ def glue_cleanup(crawler_name: str, job_name: str, db_name: str) -> None:
         # TEST TEARDOWN
         glue_cleanup(glue_crawler_name, glue_job_name, glue_db_name),
         delete_bucket,
-        delete_logs(cast(str, submit_glue_job.output), glue_crawler_name),
+        delete_logs(submit_glue_job.output, glue_crawler_name),
     )
 
     from tests.system.utils.watcher import watcher

tests/system/providers/amazon/aws/example_step_functions.py

@@ -16,7 +16,6 @@
 # under the License.
 import json
 from datetime import datetime
-from typing import cast
 
 from airflow import DAG
 from airflow.decorators import task
@@ -80,11 +79,11 @@ def delete_state_machine(state_machine_arn):
 
     # [START howto_operator_step_function_start_execution]
     start_execution = StepFunctionStartExecutionOperator(
-        task_id='start_execution', state_machine_arn=cast(str, state_machine_arn)
+        task_id='start_execution', state_machine_arn=state_machine_arn
     )
     # [END howto_operator_step_function_start_execution]
 
-    execution_arn = cast(str, start_execution.output)
+    execution_arn = start_execution.output
 
     # [START howto_sensor_step_function_execution]
     wait_for_execution = StepFunctionExecutionSensor(

tests/system/providers/dbt/cloud/example_dbt_cloud.py

@@ -16,7 +16,6 @@
 # under the License.
 
 from datetime import datetime
-from typing import cast
 
 from airflow.models import DAG
 
@@ -57,7 +56,7 @@
 
     # [START howto_operator_dbt_cloud_get_artifact]
     get_run_results_artifact = DbtCloudGetJobRunArtifactOperator(
-        task_id="get_run_results_artifact", run_id=cast(int, trigger_job_run1.output), path="run_results.json"
+        task_id="get_run_results_artifact", run_id=trigger_job_run1.output, path="run_results.json"
     )
     # [END howto_operator_dbt_cloud_get_artifact]
 
@@ -72,7 +71,7 @@
 
     # [START howto_operator_dbt_cloud_run_job_sensor]
     job_run_sensor = DbtCloudJobRunSensor(
-        task_id="job_run_sensor", run_id=cast(int, trigger_job_run2.output), timeout=20
+        task_id="job_run_sensor", run_id=trigger_job_run2.output, timeout=20
     )
     # [END howto_operator_dbt_cloud_run_job_sensor]
 

tests/system/providers/google/cloud/dataproc/example_dataproc_spark_async.py

@@ -21,7 +21,6 @@
 
 import os
 from datetime import datetime
-from typing import cast
 
 from airflow import models
 from airflow.providers.google.cloud.operators.dataproc import (
@@ -91,7 +90,7 @@
         task_id='spark_task_async_sensor_task',
         region=REGION,
         project_id=PROJECT_ID,
-        dataproc_job_id=cast(str, spark_task_async.output),
+        dataproc_job_id=spark_task_async.output,
         poke_interval=10,
     )
     # [END cloud_dataproc_async_submit_sensor]