jason810496 commented on code in PR #74634:
URL: https://github.com/apache/airflow/pull/74634#discussion_r4439007731
##########
task-sdk/src/airflow/sdk/importers/yaml_importer/models.py:
##########
@@ -133,6 +151,14 @@ def __get_pydantic_json_schema__(
schema["oneOf"] = [{"required": [key]} for key in cls._marker_keys]
return schema
+ def iter_dependency_task_ids(self) -> Iterator[str]:
+ """Task ids this node references via XCom; overridden by
:class:`XComRef`."""
+ return iter(())
+
+ def get_run_binding_error(self, *, top_level: bool = True) -> str | None:
+ """Why this node can't cross a run: handler binding, or None;
overridden by $t/$x."""
+ return None
Review Comment:
Would it be better to define `build_marker_flags` and `as_operator_argument`
on the base class as well?
```suggestion
def get_run_binding_error(self, *, top_level: bool = True) -> str | None:
"""Why this node can't cross a run: handler binding, or None;
overridden by $t/$x."""
return None
def build_marker_flags(self) -> _MarkerFlags:
raise NotImplementedError
def as_operator_argument(self, ctx: _ArgumentContext, *, rendered: bool)
-> Any:
raise NotImplementedError
```
##########
task-sdk/src/airflow/sdk/importers/yaml_importer/models.py:
##########
@@ -44,10 +44,21 @@
)
if TYPE_CHECKING:
+ from collections.abc import Iterator
+ from typing import Protocol
+
from pydantic import GetJsonSchemaHandler
from pydantic.json_schema import JsonSchemaValue
from pydantic_core import CoreSchema
+ class _ArgumentContext(Protocol):
+ """What a value model needs from the builder to turn itself into an
operator argument."""
+
+ def build_xcom_arg(self, ref: XComRef) -> Any: ...
+ def read_file(self, path: str) -> str: ...
+ def ensure_literal(self, value: Any, *, rendered: bool) -> Any: ...
Review Comment:
Builder always returns `XComArg`, and `ensure_literal` returns the value or
`LiteralValue`. Both imports fit under `TYPE_CHECKING`. How about having
explicit union return typing instead of using Any?
```suggestion
def build_xcom_arg(self, ref: XComRef) -> XComArg: ...
def read_file(self, path: str) -> str: ...
def ensure_literal(self, value: _T, *, rendered: bool) -> _T |
LiteralValue: ...
```
##########
task-sdk/src/airflow/sdk/importers/yaml_importer/models.py:
##########
@@ -279,6 +354,42 @@ def __get_pydantic_json_schema__(
]
return schema
+ def _iter_nested_values(self) -> Iterator[Value]:
+ """Yield the nested values of a container literal; nothing for a
scalar."""
+ if isinstance(self.root, dict):
+ yield from self.root.values()
+ elif isinstance(self.root, list):
+ yield from self.root
+
+ def build_marker_flags(self) -> _MarkerFlags:
+ child_flags = [child.build_marker_flags() for child in
self._iter_nested_values()]
+ return _MarkerFlags(
+ needs_render=any(m.needs_render for m in child_flags),
+ has_marker=any(m.has_marker for m in child_flags),
+ )
+
+ def iter_dependency_task_ids(self) -> Iterator[str]:
+ for child in self._iter_nested_values():
+ yield from child.iter_dependency_task_ids()
+
+ def get_run_binding_error(self, *, top_level: bool = True) -> str | None:
+ # Anything inside a literal container is nested, so a bindable
top-level $x cannot live here.
+ return next(
+ (
+ error
+ for child in self._iter_nested_values()
+ if (error := child.get_run_binding_error(top_level=False))
+ ),
+ None,
+ )
Review Comment:
Minor nit for the readability, non-blocking.
```suggestion
for child in self._iter_nested_values():
if error := child.get_run_binding_error(top_level=False):
return error
return None
```
##########
task-sdk/src/airflow/sdk/importers/yaml_importer/models.py:
##########
@@ -144,6 +170,28 @@ class XComRef(_MarkerRef):
validation_alias=AliasChoices(*XCOM_KEYS),
)
+ def build_marker_flags(self) -> _MarkerFlags:
+ return _MarkerFlags(needs_render=True, has_marker=True)
+
+ def iter_dependency_task_ids(self) -> Iterator[str]:
+ yield self.target if isinstance(self.target, str) else self.target.task
+
+ def get_run_binding_error(self, *, top_level: bool = True) -> str | None:
+ if not isinstance(self.target, str) and self.target.key is not None:
+ return (
+ "a keyed XCom ($x with 'key') is not supported in a run: task;
"
+ "only a task's return value crosses the handler boundary"
+ )
+ if not top_level:
+ return (
+ "an XCom ($x) nested inside a run argument is not supported;
pass "
+ "the upstream output as its own top-level run argument"
+ )
+ return None
+
+ def as_operator_argument(self, ctx: _ArgumentContext, *, rendered: bool)
-> Any:
Review Comment:
Every override has an exact type: `XComArg` here, `str` for `$t`, `str |
LiteralValue` for `$f`, the recursive union for a literal. Only `$const`
genuinely stays `Any`.
```suggestion
def as_operator_argument(self, ctx: _ArgumentContext, *, rendered: bool)
-> XComArg:
```
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
