This is an automated email from the ASF dual-hosted git repository.
uranusjr pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/main by this push:
new c83c0e341f Two typing fixes (#25690)
c83c0e341f is described below
commit c83c0e341f8d737e11127e285becd99c1201846c
Author: Tzu-ping Chung <[email protected]>
AuthorDate: Fri Aug 12 22:42:31 2022 +0800
Two typing fixes (#25690)
---
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 bae9322ef7..bf9a32bfee 100644
--- a/airflow/models/abstractoperator.py
+++ b/airflow/models/abstractoperator.py
@@ -27,14 +27,12 @@ from typing import (
FrozenSet,
Iterable,
List,
- MutableMapping,
Optional,
Sequence,
Set,
Tuple,
Type,
Union,
- cast,
)
from airflow.compat.functools import cached_property
@@ -413,8 +411,8 @@ class AbstractOperator(LoggingMixin, DAGNode):
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 ebdf07a382..502c863876 100644
--- a/airflow/models/taskinstance.py
+++ b/airflow/models/taskinstance.py
@@ -1844,7 +1844,7 @@ class TaskInstance(Base, LoggingMixin):
@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 6003d1d1ec..5fb593c839 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 769f7986fb..b502cb5f25 100644
--- a/airflow/utils/helpers.py
+++ b/airflow/utils/helpers.py
@@ -34,10 +34,12 @@ from typing import (
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:
