feat: Add flag to not use alias for Pydantic’s json schema generation (#9)

* feat: Add flag to not use alias for json schema generation

---------

Co-authored-by: Serhii Dimchenko <[email protected]>

Author: fbrueck

.gitignore

@@ -12,3 +12,6 @@ __pycache__/
 
 # Poetry
 poetry.lock
+
+# Temporary
+tests/tmp

README.md

@@ -1,5 +1,16 @@
 # JSON Schema to AWS Glue schema converter
 
+<!-- TOC -->
+* [JSON Schema to AWS Glue schema converter](#json-schema-to-aws-glue-schema-converter)
+  * [Installation](#installation)
+  * [What?](#what)
+  * [Why?](#why)
+  * [Example](#example)
+  * [Override the type for the AWS Glue Schema](#override-the-type-for-the-aws-glue-schema)
+  * [How it works?](#how-it-works)
+  * [Future work](#future-work)
+<!-- TOC -->
+
 ## Installation
 
 ```bash
@@ -100,6 +111,46 @@ Alternatively you can run CLI with `-o` flag to set output file location:
 pydantic-glue -f example.py -c Foo -o example.json -l
 ```
 
+If your Pydantic models use field aliases, but you prefer to display the field names in the JSON schema,
+you can enable this behavior by using the `--schema-by-name` flag.
+
+Here you can find the details regarding [pydantic aliases](https://docs.pydantic.dev/latest/concepts/alias/).
+
+The following model will be converted differently with `--schema-by-name` argument.
+
+```python
+from pydantic import BaseModel, Field
+
+class A(BaseModel):
+    hey: str = Field(alias="h")
+    ho: str
+```
+
+```bash
+pydantic-glue -f tests/data/input.py -c A
+
+2025-02-01 00:08:45,046 - INFO - Generated file content:
+{
+  "//": "Generated by pydantic-glue at 2025-01-31 23:08:45.046012+00:00. DO NOT EDIT",
+  "columns": {
+    "h": "string",
+    "ho": "string"
+  }
+}
+```
+
+```bash
+ pydantic-glue -f tests/data/input.py -c A --schema-by-name
+2025-02-01 00:09:18,381 - INFO - Generated file content:
+{
+  "//": "Generated by pydantic-glue at 2025-01-31 23:09:18.380586+00:00. DO NOT EDIT",
+  "columns": {
+    "hey": "string",
+    "ho": "string"
+  }
+}
+```
+
 ## Override the type for the AWS Glue Schema
 
 Wherever there is a `type` key in the input JSON Schema, an additional key `glue_type` may be

pydantic_glue/cli.py

@@ -3,6 +3,7 @@
 import os
 import sys
 from argparse import ArgumentParser
+from dataclasses import dataclass
 from datetime import datetime, timezone
 from pathlib import Path
 
@@ -12,24 +13,50 @@
 log = logging.getLogger(__name__)
 
 
-def cli() -> None:
+@dataclass
+class Arguments:
+    module_file: str
+    class_name: str
+    output_file: str
+    log_result: bool
+    json_schema_by_alias: bool
+
+
+def parse_args(argv: list[str]) -> Arguments:
     parser = ArgumentParser()
     parser.add_argument("-f", dest="source_file", required=True, type=str, help="Path to the python file")
     parser.add_argument("-c", dest="class_name", required=True, type=str, help="Python class name")
     parser.add_argument("-o", dest="output_file", required=False, type=str, help="Path to the output file")
     parser.add_argument(
         "-l", dest="log_result", action="store_true", default=False, help="Flag if need to print result in log"
     )
-    args = parser.parse_args(sys.argv[1:])
+    parser.add_argument(
+        "--schema-by-name",
+        dest="json_schema_by_alias",
+        action="store_false",
+        default=True,
+        help="Flag to not use name for json schema generation",
+    )
+    args = parser.parse_args(argv)
+
+    return Arguments(
+        module_file=args.source_file.removesuffix(".py"),
+        class_name=args.class_name,
+        output_file=args.output_file,
+        log_result=args.log_result,
+        json_schema_by_alias=args.json_schema_by_alias,
+    )
 
-    module_file = args.source_file.removesuffix(".py")
 
-    sys.path.append(os.path.dirname(module_file))
+def cli() -> None:
+    args = parse_args(sys.argv[1:])
+    sys.path.append(os.path.dirname(args.module_file))
+    imported = __import__(os.path.basename(args.module_file))
 
-    imported = __import__(os.path.basename(module_file))
-    imported = getattr(imported, args.class_name)
-    input_schema = json.dumps(imported.model_json_schema())
+    model = getattr(imported, args.class_name)
+    input_schema = json.dumps(model.model_json_schema(by_alias=args.json_schema_by_alias))
     converted = convert(input_schema)
+
     output = json.dumps(
         {
             "//": f"Generated by pydantic-glue at {datetime.now(tz=timezone.utc)}. DO NOT EDIT",

pyproject.toml

@@ -28,6 +28,7 @@ flake8 = "^7.0.0"
 black = "^24.4.2"
 isort = "^5.13.2"
 pre-commit = "^3.7.0"
+cli-test-helpers = "^4.1.0"
 
 [tool.autoflake]
 recursive = true

tests/data/expected.json

@@ -0,0 +1,7 @@
+{
+  "//": "Generated by pydantic-glue at 2025-01-31 22:51:59.027907+00:00. DO NOT EDIT",
+  "columns": {
+    "h": "string",
+    "ho": "string"
+  }
+}

tests/data/expected_by_name.json

@@ -0,0 +1,7 @@
+{
+  "//": "Generated by pydantic-glue at 2025-01-31 22:51:59.027907+00:00. DO NOT EDIT",
+  "columns": {
+    "hey": "string",
+    "ho": "string"
+  }
+}

tests/data/input.py

@@ -0,0 +1,6 @@
+from pydantic import BaseModel, Field
+
+
+class A(BaseModel):
+    hey: str = Field(alias="h")
+    ho: str

tests/unit/test_cli.py

@@ -0,0 +1,40 @@
+import json
+from pathlib import Path
+from unittest import TestCase
+
+from cli_test_helpers import shell
+from pydantic_glue.cli import parse_args
+
+
+def test_parse_args():
+
+    args = parse_args(["-f", "foo/bar/file.py", "-c", "Test"])
+    assert args.module_file == "foo/bar/file"
+    assert args.class_name == "Test"
+    assert args.output_file is None
+    assert args.log_result is False
+    assert args.json_schema_by_alias is True
+
+
+def test_parse_args_schema_by_name():
+
+    args = parse_args(["-f", "foo/bar/file.py", "-c", "Test", "--schema-by-name"])
+    assert args.module_file == "foo/bar/file"
+    assert args.class_name == "Test"
+    assert args.output_file is None
+    assert args.log_result is False
+    assert args.json_schema_by_alias is False
+
+
+def test_cli():
+    shell("pydantic-glue -f tests/data/input.py -c A -o tests/tmp/actual.json")
+    actual = json.loads(Path("tests/tmp/actual.json").read_text())
+    expected = json.loads(Path("tests/data/expected.json").read_text())
+    TestCase().assertDictEqual(actual["columns"], expected["columns"])
+
+
+def test_cli_schema_by_name():
+    shell("pydantic-glue -f tests/data/input.py -c A -o tests/tmp/actual_by_name.json --schema-by-name")
+    actual = json.loads(Path("tests/tmp/actual_by_name.json").read_text())
+    expected = json.loads(Path("tests/data/expected_by_name.json").read_text())
+    TestCase().assertDictEqual(actual["columns"], expected["columns"])