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

Enhance magic methods on XComArg for UX #21882

Merged
merged 2 commits into from
Mar 3, 2022
Merged
Changes from 1 commit
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
18 changes: 17 additions & 1 deletion airflow/models/xcom_arg.py
Original file line number Diff line number Diff line change
Expand Up @@ -66,10 +66,26 @@ def __init__(self, operator: "Operator", key: str = XCOM_RETURN_KEY):
def __eq__(self, other):
return self.operator == other.operator and self.key == other.key

def __getitem__(self, item):
def __getitem__(self, item: str) -> "XComArg":
"""Implements xcomresult['some_result_key']"""
if not isinstance(item, str):
raise ValueError("XComArg only supports str lookup")
return XComArg(operator=self.operator, key=item)

def __iter__(self):
"""Override iterable protocol to raise error explicitly.

The default ``__iter__`` implementation in Python calls ``__getitem__``
with 0, 1, 2, etc. until it hits an ``IndexError``. This does not work
well with our custom ``__getitem__`` implementation, and results in poor
DAG-writing experience since a misplaced ``*`` expansion would create an
infinite loop consuming the entire DAG parser.

This override catches the error eagerly, so an incorrectly implemented
DAG fails fast and avoids wasting resources on nonsensical iterating.
"""
raise TypeError(f"{self.__class__.__name__!r} object is not iterable")

def __str__(self):
"""
Backward compatibility for old-style jinja used in Airflow Operators
Expand Down