|
1 | | -from typing import ClassVar, List, Set, Union |
| 1 | +from typing import ClassVar, List, NamedTuple, Optional, Set, Tuple, Union |
2 | 2 |
|
3 | 3 | import attr |
4 | 4 |
|
| 5 | +from ... import schema as oai |
| 6 | +from ... import utils |
| 7 | +from ..errors import PropertyError |
5 | 8 | from ..reference import Reference |
6 | 9 | from .property import Property |
| 10 | +from .schemas import Schemas |
7 | 11 |
|
8 | 12 |
|
9 | 13 | @attr.s(auto_attribs=True, frozen=True) |
@@ -48,3 +52,122 @@ def get_imports(self, *, prefix: str) -> Set[str]: |
48 | 52 | } |
49 | 53 | ) |
50 | 54 | return imports |
| 55 | + |
| 56 | + |
| 57 | +class _PropertyData(NamedTuple): |
| 58 | + optional_props: List[Property] |
| 59 | + required_props: List[Property] |
| 60 | + relative_imports: Set[str] |
| 61 | + schemas: Schemas |
| 62 | + |
| 63 | + |
| 64 | +def _process_properties(*, data: oai.Schema, schemas: Schemas, class_name: str) -> Union[_PropertyData, PropertyError]: |
| 65 | + from . import property_from_data |
| 66 | + |
| 67 | + required_properties: List[Property] = [] |
| 68 | + optional_properties: List[Property] = [] |
| 69 | + relative_imports: Set[str] = set() |
| 70 | + required_set = set(data.required or []) |
| 71 | + |
| 72 | + all_props = data.properties or {} |
| 73 | + for sub_prop in data.allOf or []: |
| 74 | + if isinstance(sub_prop, oai.Reference): |
| 75 | + source_name = Reference.from_ref(sub_prop.ref).class_name |
| 76 | + sub_model = schemas.models.get(source_name) |
| 77 | + if sub_model is None: |
| 78 | + return PropertyError(f"Reference {sub_prop.ref} not found") |
| 79 | + required_properties.extend(sub_model.required_properties) |
| 80 | + optional_properties.extend(sub_model.optional_properties) |
| 81 | + relative_imports.update(sub_model.relative_imports) |
| 82 | + else: |
| 83 | + all_props.update(sub_prop.properties or {}) |
| 84 | + required_set.update(sub_prop.required or []) |
| 85 | + |
| 86 | + for key, value in all_props.items(): |
| 87 | + prop_required = key in required_set |
| 88 | + prop, schemas = property_from_data( |
| 89 | + name=key, required=prop_required, data=value, schemas=schemas, parent_name=class_name |
| 90 | + ) |
| 91 | + if isinstance(prop, PropertyError): |
| 92 | + return prop |
| 93 | + if prop_required and not prop.nullable: |
| 94 | + required_properties.append(prop) |
| 95 | + else: |
| 96 | + optional_properties.append(prop) |
| 97 | + relative_imports.update(prop.get_imports(prefix="..")) |
| 98 | + |
| 99 | + return _PropertyData( |
| 100 | + optional_props=optional_properties, |
| 101 | + required_props=required_properties, |
| 102 | + relative_imports=relative_imports, |
| 103 | + schemas=schemas, |
| 104 | + ) |
| 105 | + |
| 106 | + |
| 107 | +def build_model_property( |
| 108 | + *, data: oai.Schema, name: str, schemas: Schemas, required: bool, parent_name: Optional[str] |
| 109 | +) -> Tuple[Union[ModelProperty, PropertyError], Schemas]: |
| 110 | + """ |
| 111 | + A single ModelProperty from its OAI data |
| 112 | +
|
| 113 | + Args: |
| 114 | + data: Data of a single Schema |
| 115 | + name: Name by which the schema is referenced, such as a model name. |
| 116 | + Used to infer the type name if a `title` property is not available. |
| 117 | + schemas: Existing Schemas which have already been processed (to check name conflicts) |
| 118 | + required: Whether or not this property is required by the parent (affects typing) |
| 119 | + parent_name: The name of the property that this property is inside of (affects class naming) |
| 120 | + """ |
| 121 | + from . import property_from_data |
| 122 | + |
| 123 | + class_name = data.title or name |
| 124 | + if parent_name: |
| 125 | + class_name = f"{utils.pascal_case(parent_name)}{utils.pascal_case(class_name)}" |
| 126 | + ref = Reference.from_ref(class_name) |
| 127 | + |
| 128 | + property_data = _process_properties(data=data, schemas=schemas, class_name=class_name) |
| 129 | + if isinstance(property_data, PropertyError): |
| 130 | + return property_data, schemas |
| 131 | + schemas = property_data.schemas |
| 132 | + |
| 133 | + additional_properties: Union[bool, Property, PropertyError] |
| 134 | + if data.additionalProperties is None: |
| 135 | + additional_properties = True |
| 136 | + elif isinstance(data.additionalProperties, bool): |
| 137 | + additional_properties = data.additionalProperties |
| 138 | + elif isinstance(data.additionalProperties, oai.Schema) and not any(data.additionalProperties.dict().values()): |
| 139 | + # An empty schema |
| 140 | + additional_properties = True |
| 141 | + else: |
| 142 | + assert isinstance(data.additionalProperties, (oai.Schema, oai.Reference)) |
| 143 | + additional_properties, schemas = property_from_data( |
| 144 | + name="AdditionalProperty", |
| 145 | + required=True, # in the sense that if present in the dict will not be None |
| 146 | + data=data.additionalProperties, |
| 147 | + schemas=schemas, |
| 148 | + parent_name=class_name, |
| 149 | + ) |
| 150 | + if isinstance(additional_properties, PropertyError): |
| 151 | + return additional_properties, schemas |
| 152 | + property_data.relative_imports.update(additional_properties.get_imports(prefix="..")) |
| 153 | + |
| 154 | + prop = ModelProperty( |
| 155 | + reference=ref, |
| 156 | + required_properties=property_data.required_props, |
| 157 | + optional_properties=property_data.optional_props, |
| 158 | + relative_imports=property_data.relative_imports, |
| 159 | + description=data.description or "", |
| 160 | + default=None, |
| 161 | + nullable=data.nullable, |
| 162 | + required=required, |
| 163 | + name=name, |
| 164 | + additional_properties=additional_properties, |
| 165 | + ) |
| 166 | + if prop.reference.class_name in schemas.models: |
| 167 | + error = PropertyError( |
| 168 | + data=data, detail=f'Attempted to generate duplicate models with name "{prop.reference.class_name}"' |
| 169 | + ) |
| 170 | + return error, schemas |
| 171 | + |
| 172 | + schemas = attr.evolve(schemas, models={**schemas.models, prop.reference.class_name: prop}) |
| 173 | + return prop, schemas |
0 commit comments