Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Two typing fixes #25690

Merged
merged 1 commit into from Aug 12, 2022
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Jump to
Jump to file
Failed to load files.
Diff view
Diff view
6 changes: 2 additions & 4 deletions airflow/models/abstractoperator.py
Expand Up @@ -27,14 +27,12 @@
FrozenSet,
Iterable,
List,
MutableMapping,
Optional,
Sequence,
Set,
Tuple,
Type,
Union,
cast,
)

from airflow.compat.functools import cached_property
Expand Down Expand Up @@ -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)
Expand Down
2 changes: 1 addition & 1 deletion airflow/models/taskinstance.py
Expand Up @@ -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,
Expand Down
2 changes: 1 addition & 1 deletion airflow/utils/context.pyi
Expand Up @@ -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
Expand Down
10 changes: 6 additions & 4 deletions airflow/utils/helpers.py
Expand Up @@ -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

Expand Down Expand Up @@ -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:
Expand Down