From 355134f51dcc65addb429c03f7e210f1d9a2bb7e Mon Sep 17 00:00:00 2001 From: Ihar Hrachyshka Date: Thu, 20 Mar 2025 12:54:02 -0400 Subject: [PATCH] fix: Support types.UnionType in schemas (#1721) # What does this PR do? Since Python 3.10, unions can be expressed as `type1 | type2`. Sadly, while this is functionally equivalent to `Union[type1, type2]`, the type of the expression is different (`types.UnionType`, not `typing.Union`). We should handle both in schemas. ## Test Plan Switch a schema type from Union to `|` and confirm the generator doesn't crash with: ``` Traceback (most recent call last): File "", line 198, in _run_module_as_main File "", line 88, in _run_code File "/Users/ihrachys/src/llama-stack/docs/openapi_generator/generate.py", line 91, in fire.Fire(main) File "/Users/ihrachys/.cache/uv/archive-v0/FBgkcwcN-PaJ0NAur__7J/lib/python3.11/site-packages/fire/core.py", line 135, in Fire component_trace = _Fire(component, args, parsed_flag_args, context, name) ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ File "/Users/ihrachys/.cache/uv/archive-v0/FBgkcwcN-PaJ0NAur__7J/lib/python3.11/site-packages/fire/core.py", line 468, in _Fire component, remaining_args = _CallAndUpdateTrace( ^^^^^^^^^^^^^^^^^^^^ File "/Users/ihrachys/.cache/uv/archive-v0/FBgkcwcN-PaJ0NAur__7J/lib/python3.11/site-packages/fire/core.py", line 684, in _CallAndUpdateTrace component = fn(*varargs, **kwargs) ^^^^^^^^^^^^^^^^^^^^^^ File "/Users/ihrachys/src/llama-stack/docs/openapi_generator/generate.py", line 55, in main spec = Specification( ^^^^^^^^^^^^^^ File "/Users/ihrachys/src/llama-stack/docs/openapi_generator/pyopenapi/utility.py", line 30, in __init__ self.document = generator.generate() ^^^^^^^^^^^^^^^^^^^^ File "/Users/ihrachys/src/llama-stack/docs/openapi_generator/pyopenapi/generator.py", line 782, in generate operation = self._build_operation(op) ^^^^^^^^^^^^^^^^^^^^^^^^^ File "/Users/ihrachys/src/llama-stack/docs/openapi_generator/pyopenapi/generator.py", line 648, in _build_operation "application/json": builder.build_media_type( ^^^^^^^^^^^^^^^^^^^^^^^^^ File "/Users/ihrachys/src/llama-stack/docs/openapi_generator/pyopenapi/generator.py", line 221, in build_media_type schema = self.schema_builder.classdef_to_ref(item_type) ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ File "/Users/ihrachys/src/llama-stack/docs/openapi_generator/pyopenapi/generator.py", line 135, in classdef_to_ref type_schema = self.classdef_to_schema(typ) ^^^^^^^^^^^^^^^^^^^^^^^^^^^^ File "/Users/ihrachys/src/llama-stack/docs/openapi_generator/pyopenapi/generator.py", line 116, in classdef_to_schema type_schema, type_definitions = self.schema_generator.classdef_to_schema(typ) ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ File "/Users/ihrachys/src/llama-stack/llama_stack/strong_typing/schema.py", line 607, in classdef_to_schema types_defined[sub_name] = self._type_to_schema_with_lookup(sub_type) ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ File "/Users/ihrachys/src/llama-stack/llama_stack/strong_typing/schema.py", line 564, in _type_to_schema_with_lookup type_schema = self.type_to_schema(data_type, force_expand=True) ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ File "/Users/ihrachys/src/llama-stack/llama_stack/strong_typing/schema.py", line 320, in type_to_schema return self._type_to_schema(data_type, force_expand, json_schema_extra) | common_info ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ File "/Users/ihrachys/src/llama-stack/llama_stack/strong_typing/schema.py", line 487, in _type_to_schema property_docstrings = get_class_property_docstrings(typ, self.options.property_description_fun) ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ File "/Users/ihrachys/src/llama-stack/llama_stack/strong_typing/schema.py", line 94, in get_class_property_docstrings for base in inspect.getmro(data_type): ^^^^^^^^^^^^^^^^^^^^^^^^^ File "/nix/store/w2wykgpkzidnnr6cpw8wf94ghb0p8big-python3-3.11.11/lib/python3.11/inspect.py", line 731, in getmro return cls.__mro__ ^^^^^^^^^^^ AttributeError: 'types.UnionType' object has no attribute '__mro__'. Did you mean: '__or__'? ``` Signed-off-by: Ihar Hrachyshka --- llama_stack/strong_typing/schema.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/llama_stack/strong_typing/schema.py b/llama_stack/strong_typing/schema.py index dfc51ea78..de69c9b82 100644 --- a/llama_stack/strong_typing/schema.py +++ b/llama_stack/strong_typing/schema.py @@ -17,6 +17,7 @@ import enum import functools import inspect import json +import types import typing import uuid from copy import deepcopy @@ -455,7 +456,7 @@ class JsonSchemaGenerator: "maxItems": len(args), "prefixItems": [self.type_to_schema(member_type) for member_type in args], } - elif origin_type is Union: + elif origin_type in (Union, types.UnionType): discriminator = None if typing.get_origin(data_type) is Annotated: discriminator = typing.get_args(data_type)[1].discriminator