From f2ba5fa0034bd3809ac828d931c64095007c5173 Mon Sep 17 00:00:00 2001 From: Leon Haffmans <49658102+lord-haffi@users.noreply.github.com> Date: Fri, 10 Nov 2023 20:47:53 +0100 Subject: [PATCH] Support Pydantic v1 output (#12) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * Support Pydantic v1 output * 🩹 --- src/bo4e_generator/__main__.py | 16 ++++++++++++---- src/bo4e_generator/parser.py | 10 +++++++--- unittests/test_main.py | 2 +- 3 files changed, 20 insertions(+), 8 deletions(-) diff --git a/src/bo4e_generator/__main__.py b/src/bo4e_generator/__main__.py index 14f0217..c238fb2 100644 --- a/src/bo4e_generator/__main__.py +++ b/src/bo4e_generator/__main__.py @@ -20,14 +20,14 @@ def resolve_paths(input_directory: Path, output_directory: Path) -> tuple[Path, return input_directory, output_directory -def generate_bo4e_schemas(input_directory: Path, output_directory: Path): +def generate_bo4e_schemas(input_directory: Path, output_directory: Path, pydantic_v1: bool = False): """ Generate all BO4E schemas from the given input directory and save them in the given output directory. """ input_directory, output_directory = resolve_paths(input_directory, output_directory) namespace = get_namespace(input_directory, output_directory) for schema_metadata in namespace.values(): - result = generate_bo4e_schema(schema_metadata, namespace) + result = generate_bo4e_schema(schema_metadata, namespace, pydantic_v1) schema_metadata.save(result) print(f"Generated {schema_metadata}") create_init_files(output_directory, get_version(namespace)) @@ -49,13 +49,21 @@ def generate_bo4e_schemas(input_directory: Path, output_directory: Path): help="Output directory for the generated python files.", required=True, ) +@click.option( + "--pydantic-v1/--pydantic-v2", + "-p1/-p2", + is_flag=True, + help="Generate pydantic v1 models instead of pydantic v2 models.", + required=False, + default=False, +) @click.help_option() @click.version_option(package_name="BO4E-Python-Generator") -def main(input_dir: Path, output_dir: Path): +def main(input_dir: Path, output_dir: Path, pydantic_v1: bool): """ CLI entry point for the bo4e-generator. """ - generate_bo4e_schemas(input_dir, output_dir) + generate_bo4e_schemas(input_dir, output_dir, pydantic_v1) if __name__ == "__main__": diff --git a/src/bo4e_generator/parser.py b/src/bo4e_generator/parser.py index 36dda0a..d1889a7 100644 --- a/src/bo4e_generator/parser.py +++ b/src/bo4e_generator/parser.py @@ -193,12 +193,16 @@ def remove_future_import(python_code: str) -> str: return re.sub(r"from __future__ import annotations\n\n", "", python_code) -def generate_bo4e_schema(schema_metadata: SchemaMetadata, namespace: dict[str, SchemaMetadata]) -> str: +def generate_bo4e_schema( + schema_metadata: SchemaMetadata, namespace: dict[str, SchemaMetadata], pydantic_v1: bool = False +) -> str: """ Generate a pydantic v2 model from the given schema. Returns the resulting code as string. """ data_model_types = get_bo4e_data_model_types( - DataModelType.PydanticV2BaseModel, target_python_version=PythonVersion.PY_311, namespace=namespace + DataModelType.PydanticBaseModel if pydantic_v1 else DataModelType.PydanticV2BaseModel, + target_python_version=PythonVersion.PY_311, + namespace=namespace, ) monkey_patch_field_name_resolver() monkey_patch_relative_import() @@ -210,7 +214,7 @@ def generate_bo4e_schema(schema_metadata: SchemaMetadata, namespace: dict[str, S data_model_field_type=data_model_types.field_model, data_type_manager_type=data_model_types.data_type_manager, dump_resolve_reference_action=data_model_types.dump_resolve_reference_action, - use_annotated=True, + use_annotated=not pydantic_v1, use_double_quotes=True, use_schema_description=True, use_subclass_enum=True, diff --git a/unittests/test_main.py b/unittests/test_main.py index 4a9b210..c907a62 100644 --- a/unittests/test_main.py +++ b/unittests/test_main.py @@ -18,7 +18,7 @@ class TestMain: def test_main(self): os.chdir(BASE_DIR) runner = CliRunner() - result = runner.invoke(main, ["--input-dir", str(INPUT_DIR), "--output-dir", str(OUTPUT_DIR)]) + result = runner.invoke(main, ["--input-dir", str(INPUT_DIR), "--output-dir", str(OUTPUT_DIR), "-p2"]) assert ( result.exit_code == 0