koxudaxi / fastapi-code-generator

This code generator creates FastAPI app from an openapi file.
MIT License
1.08k stars 109 forks source link

Support for the route body type like Union[TypeA, TypeB] #391

Open alexeynick opened 1 year ago

alexeynick commented 1 year ago

I've noticed that routes with body argument be a union of types (Union[TypeA, TypeB] for example) when combined in a main.py file lack TypeB in import section (if it's not used anywhere else). Got an idea how to fix it (in visitor/imports.py):

def _get_most_of_reference(data_type: DataType) -> Iterable[Reference]:
    if data_type.reference:
        yield data_type.reference
    for data_type in data_type.data_types:
        for reference in _get_most_of_reference(data_type):
            yield reference

def get_imports(parser: OpenAPIParser, model_path: Path) -> Dict[str, object]:
    imports = Imports()

    imports.update(parser.imports)
    for data_type in parser.data_types:
        for reference in _get_most_of_reference(data_type):
            imports.append(data_type.all_imports)
            imports.append(
                Import.from_full_path(f'.{model_path.stem}.{reference.name}')
            )
    for from_, imports_ in parser.imports_for_fastapi.items():
        imports[from_].update(imports_)
    return {'imports': imports}

If we replace "return" with "yield" in _get_most_of_reference we'll be able to iterate through all the nested data types (otherwise it returns only TypeA). Thank you for your project, it's great!