jason810496 commented on code in PR #69757:
URL: https://github.com/apache/airflow/pull/69757#discussion_r3643024932
##########
providers/standard/pyproject.toml:
##########
@@ -61,6 +61,8 @@ requires-python = ">=3.10"
dependencies = [
"apache-airflow>=2.11.0",
"apache-airflow-providers-common-compat>=1.14.1", # use next version
+ # The stub decorator generates arg-binding value schemas from parameter
annotations
+ "pydantic>=2.11.0",
Review Comment:
`pydantic` is already part of the dependency in Airflow. I don't think we
need to add them explicitly here (which might also cause further dependency
conflict)
##########
airflow-core/src/airflow/api_fastapi/execution_api/datamodels/task_arg_binding.py:
##########
@@ -32,36 +32,21 @@
from airflow.api_fastapi.core_api.base import BaseModel
-# A named alias (like TaskArgBinding below) so both union branches of
ArgValueSchema.type
-# reference one shared schema definition instead of two inlined enum copies;
the explicit
-# title lets the supervisor-schema dump merge this def with the
task-sdk-generated twin.
-JsonSchemaType = TypeAliasType(
- "JsonSchemaType",
- Annotated[
- Literal["string", "integer", "number", "boolean", "object", "array",
"null"],
- Field(title="JsonSchemaType"),
- ],
+# A named alias (like TaskArgBinding below) so every schema keeps one shared,
titled
Review Comment:
Shorting all these comments in this module we introduced as they're too
verbose.
##########
providers/standard/tests/unit/standard/decorators/test_stub.py:
##########
@@ -282,51 +282,91 @@ def group(n):
pytest.param(bool, {"type": "boolean"}, id="bool"),
pytest.param(int, {"type": "integer", "format": "int64"}, id="int"),
pytest.param(float, {"type": "number", "format": "double"},
id="float"),
- pytest.param(dict, {"type": "object"}, id="dict"),
- pytest.param(dict[str, int], {"type": "object"},
id="dict-parameterized"),
- pytest.param(typing.Mapping[str, int], {"type": "object"},
id="mapping"),
- pytest.param(list, {"type": "array"}, id="list"),
- pytest.param(list[int], {"type": "array"}, id="list-parameterized"),
- pytest.param(tuple, {"type": "array"}, id="tuple"),
- pytest.param(set, {"type": "array"}, id="set"),
- pytest.param(typing.Sequence[int], {"type": "array"}, id="sequence"),
+ pytest.param(dict, {"type": "object", "additionalProperties": True},
id="dict"),
+ pytest.param(
+ dict[str, int],
+ {"type": "object", "additionalProperties": {"type": "integer",
"format": "int64"}},
+ id="dict-parameterized",
+ ),
+ pytest.param(
+ typing.Mapping[str, int],
+ {"type": "object", "additionalProperties": {"type": "integer",
"format": "int64"}},
+ id="mapping",
+ ),
+ pytest.param(list, {"type": "array", "items": {}}, id="list"),
+ pytest.param(
+ list[int],
+ {"type": "array", "items": {"type": "integer", "format": "int64"}},
+ id="list-parameterized",
+ ),
+ pytest.param(tuple, {"type": "array", "items": {}}, id="tuple"),
+ pytest.param(set, {"type": "array", "items": {}, "uniqueItems": True},
id="set"),
+ pytest.param(
+ typing.Sequence[int],
+ {"type": "array", "items": {"type": "integer", "format": "int64"}},
+ id="sequence",
+ ),
pytest.param(datetime.datetime, {"type": "string", "format":
"date-time"}, id="datetime"),
- pytest.param(pendulum.DateTime, {"type": "string", "format":
"date-time"}, id="pendulum-datetime"),
pytest.param(datetime.date, {"type": "string", "format": "date"},
id="date"),
pytest.param(datetime.time, {"type": "string", "format": "time"},
id="time"),
pytest.param(datetime.timedelta, {"type": "string", "format":
"duration"}, id="timedelta"),
+ pytest.param(bytes, {"type": "string", "format": "binary"},
id="bytes"),
+ pytest.param(
+ typing.Literal["a", "b"],
+ {"type": "string", "enum": ["a", "b"]},
+ id="literal",
+ ),
pytest.param(Any, None, id="any"),
pytest.param(None, None, id="none"),
pytest.param(type(None), None, id="nonetype"),
- pytest.param(bytes, None, id="bytes"),
+ pytest.param(
+ pendulum.DateTime,
+ None,
+ id="pendulum-datetime",
+ # pydantic has no schema for arbitrary datetime subclasses;
decode-only fallback.
+ ),
Review Comment:
The _ValueSchemaGenerator should handle the `pendulum.DateTime` as it's
really common in Airflow.
##########
providers/standard/src/airflow/providers/standard/decorators/stub.py:
##########
@@ -39,84 +40,46 @@
from airflow.providers.common.compat.sdk import Context
-def _json_schema_fragment(annotation: Any) -> dict[str, Any] | None:
- """Map one non-union annotation to a JSON-schema fragment; ``None`` =
unclassifiable."""
- if annotation is type(None):
- return {"type": "null"}
- origin = typing.get_origin(annotation)
- if origin is not None:
- annotation = origin
- if not isinstance(annotation, type):
- return None
- # bool subclasses int, str/bytes are Sequences, and datetime subclasses
date -- order matters.
- if issubclass(annotation, bool):
- return {"type": "boolean"}
- if issubclass(annotation, int):
- return {"type": "integer", "format": "int64"}
- if issubclass(annotation, float):
- return {"type": "number", "format": "double"}
- if issubclass(annotation, str):
- return {"type": "string"}
- if issubclass(annotation, bytes):
- return None
- if issubclass(annotation, datetime.datetime):
- return {"type": "string", "format": "date-time"}
- if issubclass(annotation, datetime.date):
- return {"type": "string", "format": "date"}
- if issubclass(annotation, datetime.time):
- return {"type": "string", "format": "time"}
- if issubclass(annotation, datetime.timedelta):
- return {"type": "string", "format": "duration"}
- if issubclass(annotation, (dict, Mapping)):
- return {"type": "object"}
- if issubclass(annotation, (list, tuple, set, frozenset, Sequence)):
- return {"type": "array"}
- return None
+class _ValueSchemaGenerator(GenerateJsonSchema):
+ """
+ Pydantic's stock JSON-schema generation plus OpenAPI's fixed-width numeric
formats.
+
+ A foreign runtime decodes numbers into machine types, which the bare
+ ``integer``/``number`` type names cannot convey; ``format`` is an
annotation per
+ JSON schema, so runtimes that don't know these names simply skip them.
+ """
+
+ def int_schema(self, schema):
+ return {**super().int_schema(schema), "format": "int64"}
+
+ def float_schema(self, schema):
+ return {**super().float_schema(schema), "format": "double"}
def _infer_value_schema(annotation: Any) -> dict[str, Any] | None:
"""
- Map a stub function parameter annotation to a JSON-schema fragment.
-
- Fragments carry the standard ``type`` keyword -- a single name, or a list
for union
- annotations (set semantics; member order follows the annotation) -- plus a
``format``
- annotation where the Python type implies a wire representation the type
name alone
- cannot (``int`` -> ``int64``, ``datetime`` -> ``date-time``, ``timedelta``
->
- ``duration``, ...). Returns ``None`` when the annotation gives no
constraint; the
- binding then omits ``value_schema`` and the foreign runtime falls back to a
- decode-only check.
+ Build the JSON-schema fragment for one stub parameter annotation, via
pydantic.
Review Comment:
Short the docstring comment.
--
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]