Use typed Context EVERYWHERE (#20565)

Part of #19891

Author: potiuk

airflow/example_dags/example_dag_decorator.py

@@ -23,6 +23,7 @@
 from airflow.decorators import dag, task
 from airflow.models.baseoperator import BaseOperator
 from airflow.operators.email import EmailOperator
+from airflow.utils.context import Context
 
 
 class GetRequestOperator(BaseOperator):
@@ -32,7 +33,7 @@ def __init__(self, *, url: str, **kwargs):
         super().__init__(**kwargs)
         self.url = url
 
-    def execute(self, context):
+    def execute(self, context: Context):
         return httpx.get(self.url).json()
 
 

airflow/example_dags/example_skip_dag.py

@@ -23,6 +23,7 @@
 from airflow import DAG
 from airflow.exceptions import AirflowSkipException
 from airflow.operators.dummy import DummyOperator
+from airflow.utils.context import Context
 from airflow.utils.trigger_rule import TriggerRule
 
 
@@ -32,7 +33,7 @@ class DummySkipOperator(DummyOperator):
 
     ui_color = '#e8b7e4'
 
-    def execute(self, context):
+    def execute(self, context: Context):
         raise AirflowSkipException
 
 

airflow/models/param.py

@@ -17,18 +17,16 @@
 import copy
 import json
 import warnings
-from typing import TYPE_CHECKING, Any, Dict, ItemsView, MutableMapping, Optional, ValuesView
+from typing import Any, Dict, ItemsView, MutableMapping, Optional, ValuesView
 
 import jsonschema
 from jsonschema import FormatChecker
 from jsonschema.exceptions import ValidationError
 
 from airflow.exceptions import AirflowException
+from airflow.utils.context import Context
 from airflow.utils.types import NOTSET, ArgNotSet
 
-if TYPE_CHECKING:
-    from airflow.utils.context import Context
-
 
 class Param:
     """
@@ -257,7 +255,7 @@ def __init__(self, current_dag, name: str, default: Optional[Any] = None):
         self._name = name
         self._default = default
 
-    def resolve(self, context: "Context") -> Any:
+    def resolve(self, context: Context) -> Any:
         """Pull DagParam value from DagRun context. This method is run during ``op.execute()``."""
         default = self._default
         if not self._default:

airflow/models/xcom_arg.py

@@ -14,17 +14,15 @@
 # KIND, either express or implied.  See the License for the
 # specific language governing permissions and limitations
 # under the License.
-from typing import TYPE_CHECKING, Any, List, Optional, Sequence, Union
+from typing import Any, List, Optional, Sequence, Union
 
 from airflow.exceptions import AirflowException
 from airflow.models.baseoperator import BaseOperator
 from airflow.models.taskmixin import DependencyMixin
 from airflow.models.xcom import XCOM_RETURN_KEY
+from airflow.utils.context import Context
 from airflow.utils.edgemodifier import EdgeModifier
 
-if TYPE_CHECKING:
-    from airflow.utils.context import Context
-
 
 class XComArg(DependencyMixin):
     """
@@ -130,7 +128,7 @@ def set_downstream(
         """Proxy to underlying operator set_downstream method. Required by TaskMixin."""
         self.operator.set_downstream(task_or_task_list, edge_modifier)
 
-    def resolve(self, context: "Context") -> Any:
+    def resolve(self, context: Context) -> Any:
         """
         Pull XCom value for the existing arg. This method is run during ``op.execute()``
         in respectable context.

airflow/operators/bash.py

@@ -22,6 +22,7 @@
 from airflow.exceptions import AirflowException, AirflowSkipException
 from airflow.hooks.subprocess import SubprocessHook
 from airflow.models import BaseOperator
+from airflow.utils.context import Context
 from airflow.utils.operator_helpers import context_to_airflow_vars
 
 
@@ -180,7 +181,7 @@ def get_env(self, context):
         env.update(airflow_context_vars)
         return env
 
-    def execute(self, context):
+    def execute(self, context: Context):
         if self.cwd is not None:
             if not os.path.exists(self.cwd):
                 raise AirflowException(f"Can not find the cwd: {self.cwd}")

airflow/operators/dummy.py

@@ -15,8 +15,8 @@
 # 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.utils.context import Context
 
 
 class DummyOperator(BaseOperator):
@@ -33,5 +33,5 @@ class DummyOperator(BaseOperator):
     def __init__(self, **kwargs) -> None:
         super().__init__(**kwargs)
 
-    def execute(self, context):
+    def execute(self, context: Context):
         pass

airflow/operators/email.py

@@ -18,6 +18,7 @@
 from typing import Any, Dict, List, Optional, Union
 
 from airflow.models import BaseOperator
+from airflow.utils.context import Context
 from airflow.utils.email import send_email
 
 
@@ -79,7 +80,7 @@ def __init__(
         self.conn_id = conn_id
         self.custom_headers = custom_headers
 
-    def execute(self, context):
+    def execute(self, context: Context):
         send_email(
             self.to,
             self.subject,

airflow/operators/generic_transfer.py

@@ -19,6 +19,7 @@
 
 from airflow.hooks.base import BaseHook
 from airflow.models import BaseOperator
+from airflow.utils.context import Context
 
 
 class GenericTransfer(BaseOperator):
@@ -72,7 +73,7 @@ def __init__(
         self.preoperator = preoperator
         self.insert_args = insert_args or {}
 
-    def execute(self, context):
+    def execute(self, context: Context):
         source_hook = BaseHook.get_hook(self.source_conn_id)
 
         self.log.info("Extracting data from %s", self.source_conn_id)

airflow/operators/subdag.py

@@ -34,6 +34,7 @@
 from airflow.models.pool import Pool
 from airflow.models.taskinstance import TaskInstance
 from airflow.sensors.base import BaseSensorOperator
+from airflow.utils.context import Context
 from airflow.utils.session import NEW_SESSION, create_session, provide_session
 from airflow.utils.state import State
 from airflow.utils.types import DagRunType
@@ -179,7 +180,7 @@ def pre_execute(self, context):
             if dag_run.state == State.FAILED:
                 self._reset_dag_run_and_task_instances(dag_run, execution_date)
 
-    def poke(self, context):
+    def poke(self, context: Context):
         execution_date = context['execution_date']
         dag_run = self._get_dagrun(execution_date=execution_date)
         return dag_run.state != State.RUNNING

airflow/providers/airbyte/operators/airbyte.py

@@ -15,11 +15,14 @@
 # KIND, either express or implied.  See the License for the
 # specific language governing permissions and limitations
 # under the License.
-from typing import Optional
+from typing import TYPE_CHECKING, Optional
 
 from airflow.models import BaseOperator
 from airflow.providers.airbyte.hooks.airbyte import AirbyteHook
 
+if TYPE_CHECKING:
+    from airflow.utils.context import Context
+
 
 class AirbyteTriggerSyncOperator(BaseOperator):
     """
@@ -68,7 +71,7 @@ def __init__(
         self.wait_seconds = wait_seconds
         self.asynchronous = asynchronous
 
-    def execute(self, context) -> None:
+    def execute(self, context: 'Context') -> None:
         """Create Airbyte Job and wait to finish"""
         hook = AirbyteHook(airbyte_conn_id=self.airbyte_conn_id, api_version=self.api_version)
         job_object = hook.submit_sync_connection(connection_id=self.connection_id)