Fix #389

CHANGELOG.md

@@ -23,6 +23,9 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - Fixes `TypeError` when writing a request without host header.
 - Add support for `Pydantic` `v2`: meaning feature parity with support for
  Pydantic v1 (generating OpenAPI Documentation).
+- Add support for `Union` types in sub-properties of request handlers input and
+ output types, for generating OpenAPI Documentation, both using simple classes
+ and Pydantic [#389](https://github.com/Neoteroi/BlackSheep/issues/389)
 
 ## [2.0a7] - 2023-05-31 :corn:
 

blacksheep/server/openapi/exceptions.py

@@ -10,13 +10,3 @@ def __init__(self, content_type: str) -> None:
  "have unique type."
  )
  self.content_type = content_type
-
-
-class UnsupportedUnionTypeException(DocumentationException):
- def __init__(self, unsupported_type) -> None:
- super().__init__(
- f"Union types are not supported for automatic generation of "
- "OpenAPI Documentation. The annotation that caused exception is: "
- f"{unsupported_type}."
- )
- self.unsupported_type = unsupported_type

blacksheep/server/openapi/v3.py

@@ -52,10 +52,7 @@
  DocstringInfo,
  get_handler_docstring_info,
 )
-from blacksheep.server.openapi.exceptions import (
- DuplicatedContentTypeDocsException,
- UnsupportedUnionTypeException,
-)
+from blacksheep.server.openapi.exceptions import DuplicatedContentTypeDocsException
 from blacksheep.server.routing import Router
 
 from ..application import Application
@@ -115,14 +112,13 @@ def check_union(object_type: Any) -> Tuple[bool, Any]:
  if hasattr(object_type, "is_required"):
  # Pydantic v2
  return object_type.is_required(), object_type
- # support only Union[None, Type] - that is equivalent of Optional[Type]
- if type(None) not in object_type.__args__ or len(object_type.__args__) > 2:
- raise UnsupportedUnionTypeException(object_type)
 
- for possible_type in object_type.__args__:
- if type(None) is possible_type:
- continue
- return True, possible_type
+ if type(None) in object_type.__args__ and len(object_type.__args__) == 2:
+ for possible_type in object_type.__args__:
+ if type(None) is possible_type:
+ continue
+ return True, possible_type
+ return type(None) not in object_type.__args__, object_type
  return False, object_type
 
 
@@ -177,10 +173,6 @@ def get_type_fields(self, object_type, register_type) -> List[FieldInfo]:
  return [FieldInfo(field.name, field.type) for field in fields(object_type)]
 
 
-def _is_optional_any_of(any_of) -> bool:
- return len(any_of) == 2 and any(item.get("type") == "null" for item in any_of)
-
-
 def _try_is_subclass(object_type, check_type):
  try:
  return issubclass(object_type, check_type)
@@ -214,9 +206,6 @@ def _allow_none(self, field_info):
  return not field_info.is_required()
 
  def _handle_any_of(self, any_of, types_args, register_type):
- # Currently only optional
- assert _is_optional_any_of(any_of)
-
  for item in any_of:
  if "$ref" in item:
  obj_type_name = item["$ref"].split("/")[-1]
@@ -591,15 +580,21 @@ def _handle_object_type(
  required: List[str],
  context_type_args: Optional[Dict[Any, Type]] = None,
  ) -> Reference:
- type_name = self.get_type_name(object_type, context_type_args)
- reference = self._register_schema(
+ return self._handle_object_type_schema(
+ object_type,
+ context_type_args,
  Schema(
  type=ValueType.OBJECT,
  required=required or None,
  properties=properties,
  ),
- type_name,
  )
+
+ def _handle_object_type_schema(
+ self, object_type, context_type_args, schema: Schema
+ ):
+ type_name = self.get_type_name(object_type, context_type_args)
+ reference = self._register_schema(schema, type_name)
  self._objects_references[object_type] = reference
  self._objects_references[type_name] = reference
  return reference
@@ -738,6 +733,18 @@ def _try_get_schema_for_generic(
  ) -> Optional[Reference]:
  origin = get_origin(object_type)
 
+ if origin is Union:
+ schema = Schema(
+ ValueType.OBJECT,
+ any_of=[
+ self.get_schema_by_type(child_type, context_type_args)
+ for child_type in object_type.__args__
+ ],
+ )
+ return self._handle_object_type_schema(
+ object_type, context_type_args, schema
+ )
+
  required: List[str] = []
  properties: Dict[str, Union[Schema, Reference]] = {}
 

tests/test_openapi_v3.py

@@ -33,10 +33,7 @@
  ResponseInfo,
  SecurityInfo,
 )
-from blacksheep.server.openapi.exceptions import (
- DuplicatedContentTypeDocsException,
- UnsupportedUnionTypeException,
-)
+from blacksheep.server.openapi.exceptions import DuplicatedContentTypeDocsException
 from blacksheep.server.openapi.v3 import (
  DataClassTypeHandler,
  OpenAPIHandler,
@@ -296,11 +293,6 @@ async def example():
  await app.start()
 
 
-def test_raises_for_union_type(docs):
- with pytest.raises(UnsupportedUnionTypeException):
- docs.get_schema_by_type(Union[Foo, Ufo])
-
-
 @pytest.mark.parametrize(
  "annotation,expected_result",
  [
@@ -3037,3 +3029,226 @@ def create_cat(self, cat: Cat) -> None:
 - name: Parrots
 """.strip()
  )
+
+
+@dataclass
+class A:
+ a_prop: int
+
+
+@dataclass
+class B:
+ b_prop: str
+
+
+@dataclass
+class C:
+ c_prop: str
+
+
+@dataclass
+class D:
+ d_prop: float
+
+
+@dataclass
+class E:
+ e_prop: int
+
+
+@dataclass
+class F:
+ f_prop: str
+ f_prop2: A
+
+
+@dataclass
+class AnyOfTestClass:
+ sub_prop: Union[A, B, C]
+
+
+@dataclass
+class AnyOfResponseTestClass:
+ data: Union[D, E, F]
+
+
+class APyd(BaseModel):
+ a_prop: int
+
+
+class BPyd(BaseModel):
+ b_prop: str
+
+
+class CPyd(BaseModel):
+ c_prop: str
+
+
+class DPyd(BaseModel):
+ d_prop: float
+
+
+class EPyd(BaseModel):
+ e_prop: int
+
+
+class FPyd(BaseModel):
+ f_prop: str
+ f_prop2: APyd
+
+
+class AnyOfTestClassPyd(BaseModel):
+ sub_prop: Union[APyd, BPyd, CPyd]
+
+
+class AnyOfResponseTestClassPyd(BaseModel):
+ data: Union[DPyd, EPyd, FPyd]
+
+
+@pytest.mark.asyncio
+async def test_any_of_dataclasses(docs: OpenAPIHandler, serializer: Serializer):
+ app = get_app()
+ docs.bind_app(app)
+
+ @app.router.post("/one")
+ def one(data: AnyOfTestClass) -> AnyOfResponseTestClass:
+ ...
+
+ await app.start()
+
+ yaml = serializer.to_yaml(docs.generate_documentation(app))
+
+ expected_fragments = [
+ """
+ /one:
+ post:
+ responses:
+ '200':
+ description: Success response
+ content:
+ application/json:
+ schema:
+ $ref: '#/components/schemas/AnyOfResponseTestClass'
+ operationId: one
+ parameters: []
+ requestBody:
+ content:
+ application/json:
+ schema:
+ $ref: '#/components/schemas/AnyOfTestClass'
+ required: true
+ """,
+ """
+ D:
+ type: object
+ required:
+ - d_prop
+ properties:
+ d_prop:
+ type: number
+ format: float
+ nullable: false
+ """,
+ """
+ UnionOfDAndEAndF:
+ type: object
+ anyOf:
+ - $ref: '#/components/schemas/D'
+ - $ref: '#/components/schemas/E'
+ - $ref: '#/components/schemas/F'
+ """,
+ """
+ AnyOfResponseTestClass:
+ type: object
+ properties:
+ data:
+ $ref: '#/components/schemas/UnionOfDAndEAndF'
+ """,
+ """
+ UnionOfAAndBAndC:
+ type: object
+ anyOf:
+ - $ref: '#/components/schemas/A'
+ - $ref: '#/components/schemas/B'
+ - $ref: '#/components/schemas/C'
+ """,
+ ]
+
+ for fragment in expected_fragments:
+ assert fragment.strip() in yaml
+
+
+@pytest.mark.asyncio
+async def test_any_of_pydantic_models(docs: OpenAPIHandler, serializer: Serializer):
+ app = get_app()
+ docs.bind_app(app)
+
+ @app.router.post("/one")
+ def one(data: AnyOfTestClassPyd) -> AnyOfResponseTestClassPyd:
+ ...
+
+ await app.start()
+
+ yaml = serializer.to_yaml(docs.generate_documentation(app))
+
+ expected_fragments = [
+ """
+ /one:
+ post:
+ responses:
+ '200':
+ description: Success response
+ content:
+ application/json:
+ schema:
+ $ref: '#/components/schemas/AnyOfResponseTestClassPyd'
+ operationId: one
+ parameters: []
+ requestBody:
+ content:
+ application/json:
+ schema:
+ $ref: '#/components/schemas/AnyOfTestClassPyd'
+ required: true
+ """,
+ """
+ DPyd:
+ type: object
+ required:
+ - d_prop
+ properties:
+ d_prop:
+ type: number
+ format: float
+ nullable: false
+ """,
+ """
+ AnyOfResponseTestClassPyd:
+ type: object
+ required:
+ - data
+ properties:
+ data:
+ title: Data
+ anyOf:
+ - $ref: '#/components/schemas/DPyd'
+ - $ref: '#/components/schemas/EPyd'
+ - $ref: '#/components/schemas/FPyd'
+ """,
+ """
+ AnyOfTestClassPyd:
+ type: object
+ required:
+ - sub_prop
+ properties:
+ sub_prop:
+ title: Sub Prop
+ anyOf:
+ - $ref: '#/components/schemas/APyd'
+ - $ref: '#/components/schemas/BPyd'
+ - $ref: '#/components/schemas/CPyd'
+ """,
+ ]
+
+ for fragment in expected_fragments:
+ assert fragment.strip() in yaml