From 227effa4fc8ceabecb4863593158037f6dd03b0e Mon Sep 17 00:00:00 2001
From: Tzu-ping Chung <uranusjr@gmail.com>
Date: Fri, 12 Aug 2022 21:44:31 +0800
Subject: [PATCH] Two typing fixes

* Correctly handle KeyboardInterrupt in TaskInstance.handle_failure()
* Move the Context-MutableMapping cast into helpers so they don't
  complain about Context not being a MutableMapping
---
 airflow/models/abstractoperator.py |  6 ++----
 airflow/models/taskinstance.py     |  2 +-
 airflow/utils/context.pyi          |  2 +-
 airflow/utils/helpers.py           | 10 ++++++----
 4 files changed, 10 insertions(+), 10 deletions(-)

diff --git a/airflow/models/abstractoperator.py b/airflow/models/abstractoperator.py
index bae9322ef7bbc..bf9a32bfeef14 100644
--- a/airflow/models/abstractoperator.py
+++ b/airflow/models/abstractoperator.py
@@ -27,14 +27,12 @@
     FrozenSet,
     Iterable,
     List,
-    MutableMapping,
     Optional,
     Sequence,
     Set,
     Tuple,
     Type,
     Union,
-    cast,
 )
 
 from airflow.compat.functools import cached_property
@@ -413,8 +411,8 @@ def render_template(
                 template = jinja_env.from_string(value)
             dag = self.get_dag()
             if dag and dag.render_template_as_native_obj:
-                return render_template_as_native(template, cast(MutableMapping[str, Any], context))
-            return render_template_to_string(template, cast(MutableMapping[str, Any], context))
+                return render_template_as_native(template, context)
+            return render_template_to_string(template, context)
 
         if isinstance(value, (DagParam, XComArg)):
             return value.resolve(context)
diff --git a/airflow/models/taskinstance.py b/airflow/models/taskinstance.py
index ebdf07a3820e1..502c863876c50 100644
--- a/airflow/models/taskinstance.py
+++ b/airflow/models/taskinstance.py
@@ -1844,7 +1844,7 @@ def get_truncated_error_traceback(error: BaseException, truncate_to: Callable) -
     @provide_session
     def handle_failure(
         self,
-        error: Union[None, str, BaseException],
+        error: Union[None, str, Exception, KeyboardInterrupt],
         test_mode: Optional[bool] = None,
         context: Optional[Context] = None,
         force_fail: bool = False,
diff --git a/airflow/utils/context.pyi b/airflow/utils/context.pyi
index 6003d1d1ec7d6..5fb593c83932e 100644
--- a/airflow/utils/context.pyi
+++ b/airflow/utils/context.pyi
@@ -61,7 +61,7 @@ class Context(TypedDict):
     ds: str
     ds_nodash: str
     execution_date: DateTime
-    exception: Union[Exception, str, None]
+    exception: Union[KeyboardInterrupt, Exception, str, None]
     inlets: list
     logical_date: DateTime
     macros: Any
diff --git a/airflow/utils/helpers.py b/airflow/utils/helpers.py
index 769f7986fbceb..b502cb5f25122 100644
--- a/airflow/utils/helpers.py
+++ b/airflow/utils/helpers.py
@@ -34,10 +34,12 @@
     Optional,
     Tuple,
     TypeVar,
+    cast,
 )
 
 from airflow.configuration import conf
 from airflow.exceptions import AirflowException
+from airflow.utils.context import Context
 from airflow.utils.module_loading import import_string
 from airflow.utils.types import NOTSET
 
@@ -292,14 +294,14 @@ def render_template(template: Any, context: MutableMapping[str, Any], *, native:
     return "".join(nodes)
 
 
-def render_template_to_string(template: "jinja2.Template", context: MutableMapping[str, Any]) -> str:
+def render_template_to_string(template: "jinja2.Template", context: Context) -> str:
     """Shorthand to ``render_template(native=False)`` with better typing support."""
-    return render_template(template, context, native=False)
+    return render_template(template, cast(MutableMapping[str, Any], context), native=False)
 
 
-def render_template_as_native(template: "jinja2.Template", context: MutableMapping[str, Any]) -> Any:
+def render_template_as_native(template: "jinja2.Template", context: Context) -> Any:
     """Shorthand to ``render_template(native=True)`` with better typing support."""
-    return render_template(template, context, native=True)
+    return render_template(template, cast(MutableMapping[str, Any], context), native=True)
 
 
 def exactly_one(*args) -> bool: