GC-117703: sdk adjustments

fastapi_code_generator/__main__.py

@@ -1,11 +1,13 @@
 import re
+import json
 from datetime import datetime, timezone
 from importlib.util import module_from_spec, spec_from_file_location
 from pathlib import Path
+from collections import defaultdict
 from typing import Any, Dict, List, Optional
 
 import typer
-from datamodel_code_generator import DataModelType, LiteralType, PythonVersion, chdir
+from datamodel_code_generator import DataModelType, LiteralType, PythonVersion, chdir, DatetimeClassType
 from datamodel_code_generator.format import CodeFormatter
 from datamodel_code_generator.imports import Import, Imports
 from datamodel_code_generator.model import get_data_model_types
@@ -65,15 +67,38 @@ def main(
     python_version: PythonVersion = typer.Option(
         PythonVersion.PY_39.value, "--python-version", "-p"
     ),
+    capitalise_enum_members: bool = typer.Option(False, "--capitalise-enum-members"),
+    output_datetime_class: Optional[DatetimeClassType] = typer.Option(
+        DatetimeClassType.Datetime, "--output-datetime-class",
+        help="Specify the datetime class to use for datetime fields"
+    ),
+    allow_population_by_field_name: bool = typer.Option(False, "--allow-population-by-field-name"),
+    extra_template_data: str = typer.Option(None, "--extra-template-data"),
+    additional_imports: str = typer.Option(None, "--additional-imports"),
 ) -> None:
     input_name: str = input_file
-    input_text: str
+    input_text: str = None
+
+    input_name = Path(input_name).expanduser().resolve()
 
-    with open(input_file, encoding=encoding) as f:
-        input_text = f.read()
+    try:
+        with open(input_file, encoding=encoding) as f:
+            input_text = f.read()
+    except:
+        pass
+
+    if extra_template_data:
+        try:
+            with open(extra_template_data, encoding=encoding) as f:
+                extra_template_data = json.load(f, object_hook=lambda d: defaultdict(dict, **d))
+        except Exception as exc:
+            print(f"could not load extra: {exc}")
 
     model_path = Path(model_file) if model_file else MODEL_PATH  # pragma: no cover
 
+    if additional_imports:
+        additional_imports = additional_imports.split(",")
+
     return generate_code(
         input_name,
         input_text,
@@ -89,6 +114,11 @@ def main(
         specify_tags=specify_tags,
         output_model_type=output_model_type,
         python_version=python_version,
+        capitalise_enum_members=capitalise_enum_members,
+        output_datetime_class=output_datetime_class,
+        allow_population_by_field_name=allow_population_by_field_name,
+        extra_template_data=extra_template_data,
+        additional_imports=additional_imports
     )
 
 
@@ -101,7 +131,6 @@ def _get_most_of_reference(data_type: DataType) -> Optional[Reference]:
             return reference
     return None
 
-
 def generate_code(
     input_name: str,
     input_text: str,
@@ -117,6 +146,11 @@ def generate_code(
     specify_tags: Optional[str] = None,
     output_model_type: DataModelType = DataModelType.PydanticBaseModel,
     python_version: PythonVersion = PythonVersion.PY_39,
+    capitalise_enum_members: bool = False,
+    output_datetime_class: Optional[DatetimeClassType] = None,
+    extra_template_data: defaultdict[str, dict[str, Any]] | None = None,
+    allow_population_by_field_name: Optional[bool] = False,
+    additional_imports: Optional[list[str]] = None
 ) -> None:
     if not model_path:
         model_path = MODEL_PATH
@@ -132,8 +166,10 @@ def generate_code(
         custom_visitors = []
     data_model_types = get_data_model_types(output_model_type, python_version)
 
+    source = input_text or input_name
+
     parser = OpenAPIParser(
-        input_text,
+        source=source,
         enum_field_as_literal=enum_field_as_literal,
         data_model_type=data_model_types.data_model,
         data_model_root_type=data_model_types.root_model,
@@ -142,6 +178,13 @@ def generate_code(
         dump_resolve_reference_action=data_model_types.dump_resolve_reference_action,
         custom_template_dir=model_template_dir,
         target_python_version=python_version,
+        additional_imports=additional_imports,
+        base_path=Path(input_name).absolute().parent,
+        capitalise_enum_members=capitalise_enum_members,
+        output_datetime_class=output_datetime_class,
+        extra_template_data=extra_template_data,
+        allow_population_by_field_name=allow_population_by_field_name,
+        field_extra_keys={"union_mode"}
     )
 
     with chdir(output_dir):
@@ -153,7 +196,7 @@ def generate_code(
         modules = {output_dir / model_path.with_suffix('.py'): (models, input_name)}
     else:
         modules = {
-            output_dir / model_path / module_name[0]: (model.body, input_name)
+            output_dir / model_path / Path(*module_name): (model.body, input_name)
             for module_name, model in models.items()
         }
 

fastapi_code_generator/parser.py

@@ -27,6 +27,7 @@
     OpenAPIScope,
     PythonVersion,
     snooper_to_methods,
+    DatetimeClassType,
 )
 from datamodel_code_generator.imports import Import, Imports
 from datamodel_code_generator.model import DataModel, DataModelFieldBase
@@ -51,7 +52,7 @@
 class CachedPropertyModel(BaseModel):
     class Config:
         arbitrary_types_allowed = True
-        keep_untouched = (cached_property,)
+        ignored_types = (cached_property,)
 
 
 class Response(BaseModel):
@@ -265,6 +266,9 @@ def __init__(
         custom_class_name_generator: Optional[Callable[[str], str]] = None,
         field_extra_keys: Optional[Set[str]] = None,
         field_include_all_keys: bool = False,
+        capitalise_enum_members: bool = False,
+        additional_imports: List[str] = None,
+        output_datetime_class: Optional[DatetimeClassType] = None,
     ):
         super().__init__(
             source=source,
@@ -282,6 +286,7 @@ def __init__(
             snake_case_field=snake_case_field,
             strip_default_none=strip_default_none,
             aliases=aliases,
+            additional_imports=additional_imports,
             allow_population_by_field_name=allow_population_by_field_name,
             apply_default_values_for_required_fields=apply_default_values_for_required_fields,
             force_optional_for_required_fields=force_optional_for_required_fields,
@@ -304,6 +309,8 @@ def __init__(
             field_extra_keys=field_extra_keys,
             field_include_all_keys=field_include_all_keys,
             openapi_scopes=[OpenAPIScope.Schemas, OpenAPIScope.Paths],
+            capitalise_enum_members=capitalise_enum_members,
+            target_datetime_class=output_datetime_class
         )
         self.operations: Dict[str, Operation] = {}
         self._temporary_operation: Dict[str, Any] = {}

fastapi_code_generator/visitors/imports.py

@@ -26,9 +26,10 @@ def get_imports(parser: OpenAPIParser, model_path: Path) -> Dict[str, object]:
     for data_type in parser.data_types:
         reference = _get_most_of_reference(data_type)
         if reference:
+            reference_path = _get_path_of_reference(model_path, reference)
             imports.append(data_type.all_imports)
             imports.append(
-                Import.from_full_path(f'.{model_path.stem}.{reference.name}')
+                Import.from_full_path(f'.{reference_path}.{reference.name}')
             )
     for from_, imports_ in parser.imports_for_fastapi.items():
         imports[from_].update(imports_)
@@ -37,5 +38,15 @@ def get_imports(parser: OpenAPIParser, model_path: Path) -> Dict[str, object]:
             imports.alias.update(operation.imports.alias)
     return {'imports': imports}
 
+def _get_path_of_reference(model_path, reference):
+    m_path = model_path.stem
+    try:
+        yaml, *_ = reference.path.split('#')
+        mod, *_ = yaml.split('.')
+        return f"{m_path}.{'.'.join(m for m in mod.split('/'))}"
+    except Exception:
+        pass
+    return m_path
+
 
 visit: Visitor = get_imports