|
4 | 4 | from enum import Enum |
5 | 5 | from typing import Any, Dict, Generator, Iterable, List, Optional, Set, Union |
6 | 6 |
|
7 | | -from .properties import EnumListProperty, EnumProperty, Property, ReferenceListProperty, RefProperty, property_from_dict |
| 7 | +from .properties import ( |
| 8 | + DateProperty, |
| 9 | + DateTimeProperty, |
| 10 | + EnumListProperty, |
| 11 | + EnumProperty, |
| 12 | + Property, |
| 13 | + ReferenceListProperty, |
| 14 | + RefProperty, |
| 15 | + property_from_dict, |
| 16 | +) |
8 | 17 | from .reference import Reference |
9 | 18 | from .responses import ListRefResponse, RefResponse, Response, response_from_dict |
10 | 19 |
|
@@ -108,11 +117,16 @@ def _add_parameters(self, data: Dict[str, Any]) -> None: |
108 | 117 | prop = property_from_dict( |
109 | 118 | name=param_dict["name"], required=param_dict["required"], data=param_dict["schema"] |
110 | 119 | ) |
111 | | - if ( |
| 120 | + if isinstance(prop, DateProperty): |
| 121 | + self.relative_imports.add("from datetime import date") |
| 122 | + elif isinstance(prop, DateTimeProperty): |
| 123 | + self.relative_imports.add("from datetime import datetime") |
| 124 | + elif ( |
112 | 125 | isinstance(prop, (ReferenceListProperty, EnumListProperty, RefProperty, EnumProperty)) |
113 | 126 | and prop.reference |
114 | 127 | ): |
115 | 128 | self.relative_imports.add(import_string_from_reference(prop.reference, prefix="..models")) |
| 129 | + |
116 | 130 | if param_dict["in"] == ParameterLocation.QUERY: |
117 | 131 | self.query_parameters.append(prop) |
118 | 132 | elif param_dict["in"] == ParameterLocation.PATH: |
@@ -166,18 +180,22 @@ def from_dict(d: Dict[str, Any], /, name: str) -> Schema: |
166 | 180 |
|
167 | 181 | ref = Reference.from_ref(d.get("title", name)) |
168 | 182 |
|
169 | | - if "properties" in d: |
170 | | - for key, value in d["properties"].items(): |
171 | | - required = key in required_set |
172 | | - p = property_from_dict(name=key, required=required, data=value) |
173 | | - if required: |
174 | | - required_properties.append(p) |
175 | | - else: |
176 | | - optional_properties.append(p) |
177 | | - if isinstance(p, (ReferenceListProperty, EnumListProperty, RefProperty, EnumProperty)) and p.reference: |
178 | | - # don't add an import for self-referencing schemas |
179 | | - if p.reference.class_name != ref.class_name: |
180 | | - relative_imports.add(import_string_from_reference(p.reference)) |
| 183 | + for key, value in d.get("properties", {}).items(): |
| 184 | + required = key in required_set |
| 185 | + p = property_from_dict(name=key, required=required, data=value) |
| 186 | + if required: |
| 187 | + required_properties.append(p) |
| 188 | + else: |
| 189 | + optional_properties.append(p) |
| 190 | + if isinstance(p, DateTimeProperty): |
| 191 | + relative_imports.add("from datetime import datetime") |
| 192 | + elif isinstance(p, DateProperty): |
| 193 | + relative_imports.add("from datetime import date") |
| 194 | + elif isinstance(p, (ReferenceListProperty, EnumListProperty, RefProperty, EnumProperty)) and p.reference: |
| 195 | + # don't add an import for self-referencing schemas |
| 196 | + if p.reference.class_name != ref.class_name: |
| 197 | + relative_imports.add(import_string_from_reference(p.reference)) |
| 198 | + |
181 | 199 | schema = Schema( |
182 | 200 | reference=ref, |
183 | 201 | required_properties=required_properties, |
|
0 commit comments