Skip to content
2 changes: 2 additions & 0 deletions docs/cli-reference/model-customization.md
Original file line number Diff line number Diff line change
Expand Up @@ -927,6 +927,8 @@ You can specify either a single base class as a string, or multiple base classes
- Single: `{"Person": "custom.bases.PersonBase"}`
- Multiple: `{"User": ["mixins.AuditMixin", "mixins.TimestampMixin"]}`

You can pass the mapping either inline as JSON or as a path to a JSON file.

When using multiple base classes, the specified classes are used directly without
adding `BaseModel`. Ensure your mixins inherit from `BaseModel` if needed.

Expand Down
2 changes: 2 additions & 0 deletions docs/cli-reference/typing-customization.md
Original file line number Diff line number Diff line change
Expand Up @@ -1535,6 +1535,8 @@ Override enum/literal generation per-field via JSON mapping.
The `--enum-field-as-literal-map` option allows per-field control over whether
to generate Literal types or Enum classes. Overrides `--enum-field-as-literal`.

You can pass the mapping either inline as JSON or as a path to a JSON file.

!!! tip "Usage"

```bash
Expand Down
4 changes: 4 additions & 0 deletions docs/llms-full.txt
Original file line number Diff line number Diff line change
Expand Up @@ -2156,6 +2156,8 @@ You can specify either a single base class as a string, or multiple base classes
- Single: `{"Person": "custom.bases.PersonBase"}`
- Multiple: `{"User": ["mixins.AuditMixin", "mixins.TimestampMixin"]}`

You can pass the mapping either inline as JSON or as a path to a JSON file.

When using multiple base classes, the specified classes are used directly without
adding `BaseModel`. Ensure your mixins inherit from `BaseModel` if needed.

Expand Down Expand Up @@ -12717,6 +12719,8 @@ Override enum/literal generation per-field via JSON mapping.
The `--enum-field-as-literal-map` option allows per-field control over whether
to generate Literal types or Enum classes. Overrides `--enum-field-as-literal`.

You can pass the mapping either inline as JSON or as a path to a JSON file.

!!! tip "Usage"

```bash
Expand Down
26 changes: 22 additions & 4 deletions src/datamodel_code_generator/arguments.py
Original file line number Diff line number Diff line change
Expand Up @@ -85,6 +85,24 @@ def _external_ref_mapping(value: str) -> str:
return value


def _json_value_or_file(value: str) -> object:
"""Parse a JSON value or load it from a JSON file path."""
path = Path(value).expanduser()
if path.is_file():
try:
json_input = path.read_text(encoding=DEFAULT_ENCODING)
except (OSError, UnicodeDecodeError) as e:
msg = f"Unable to read JSON file {value!r}: {e}"
raise ArgumentTypeError(msg) from e
else:
json_input = value
try:
return json.loads(json_input)
except json.JSONDecodeError as e:
msg = f"Invalid JSON: {e}"
raise ArgumentTypeError(msg) from e


class SortingHelpFormatter(RawDescriptionHelpFormatter):
"""Help formatter that sorts arguments, adds color to section headers, and preserves epilog formatting."""

Expand Down Expand Up @@ -505,10 +523,10 @@ def start_section(self, heading: str | None) -> None:
)
typing_options.add_argument(
"--base-class-map",
help="Model-specific base class mapping (JSON). "
help="Model-specific base class mapping (JSON or JSON file path). "
'Example: \'{"MyModel": "custom.BaseA", "OtherModel": "custom.BaseB"}\'. '
"Priority: base-class-map > customBasePath (in schema) > base-class.",
type=json.loads,
type=_json_value_or_file,
default=None,
)
typing_options.add_argument(
Expand All @@ -522,11 +540,11 @@ def start_section(self, heading: str | None) -> None:
)
typing_options.add_argument(
"--enum-field-as-literal-map",
help="Per-field override for enum/literal generation. "
help="Per-field override for enum/literal generation (JSON or JSON file path). "
"Format: JSON object mapping field names to 'literal' or 'enum'. "
'Example: \'{"status": "literal", "priority": "enum"}\'. '
"Overrides --enum-field-as-literal for matched fields.",
type=json.loads,
type=_json_value_or_file,
default=None,
)
typing_options.add_argument(
Expand Down
75 changes: 74 additions & 1 deletion tests/main/jsonschema/test_main_jsonschema.py
Original file line number Diff line number Diff line change
Expand Up @@ -2981,6 +2981,8 @@ def test_main_jsonschema_custom_base_path(output_file: Path) -> None:
- Single: `{"Person": "custom.bases.PersonBase"}`
- Multiple: `{"User": ["mixins.AuditMixin", "mixins.TimestampMixin"]}`

You can pass the mapping either inline as JSON or as a path to a JSON file.

When using multiple base classes, the specified classes are used directly without
adding `BaseModel`. Ensure your mixins inherit from `BaseModel` if needed.""",
input_schema="jsonschema/base_class_map.json",
Expand Down Expand Up @@ -3010,6 +3012,23 @@ def test_main_jsonschema_base_class_map(output_file: Path) -> None:
)


def test_main_jsonschema_base_class_map_from_file(output_file: Path, tmp_path: Path) -> None:
"""Test base_class_map loaded from a JSON file."""
mapping_path = tmp_path / "base_class_map.json"
mapping_path.write_text(
json.dumps({"Person": "custom.bases.PersonBase", "Animal": "custom.bases.AnimalBase"}),
encoding="utf-8",
)
run_main_and_assert(
input_path=JSON_SCHEMA_DATA_PATH / "base_class_map.json",
output_path=output_file,
input_file_type="jsonschema",
assert_func=assert_file_content,
expected_file="base_class_map.py",
extra_args=["--base-class-map", str(mapping_path)],
)


def test_main_jsonschema_custom_base_paths_list(output_file: Path) -> None:
"""Test customBasePath with list of base classes."""
run_main_and_assert(
Expand Down Expand Up @@ -4527,7 +4546,9 @@ def test_main_typed_dict_no_closed(output_file: Path) -> None:
option_description="""Override enum/literal generation per-field via JSON mapping.

The `--enum-field-as-literal-map` option allows per-field control over whether
to generate Literal types or Enum classes. Overrides `--enum-field-as-literal`.""",
to generate Literal types or Enum classes. Overrides `--enum-field-as-literal`.

You can pass the mapping either inline as JSON or as a path to a JSON file.""",
input_schema="jsonschema/enum_field_as_literal_map.json",
cli_args=["--enum-field-as-literal-map", '{"status": "literal"}'],
golden_output="jsonschema/enum_field_as_literal_map.py",
Expand All @@ -4551,6 +4572,20 @@ def test_main_enum_field_as_literal_map(output_file: Path) -> None:
)


def test_main_enum_field_as_literal_map_from_file(output_file: Path, tmp_path: Path) -> None:
"""Test enum_field_as_literal_map loaded from a JSON file."""
mapping_path = tmp_path / "enum_field_as_literal_map.json"
mapping_path.write_text(json.dumps({"status": "literal"}), encoding="utf-8")
run_main_and_assert(
input_path=JSON_SCHEMA_DATA_PATH / "enum_field_as_literal_map.json",
output_path=output_file,
input_file_type=None,
assert_func=assert_file_content,
expected_file="enum_field_as_literal_map.py",
extra_args=["--enum-field-as-literal-map", str(mapping_path)],
)


def test_main_enum_field_as_literal_map_override_global(output_file: Path) -> None:
"""Test --enum-field-as-literal-map overrides global --enum-field-as-literal."""
run_main_and_assert(
Expand All @@ -4568,6 +4603,44 @@ def test_main_enum_field_as_literal_map_override_global(output_file: Path) -> No
)


def test_main_enum_field_as_literal_map_invalid_json_file(tmp_path: Path, capsys: pytest.CaptureFixture[str]) -> None:
"""Test invalid JSON file passed to --enum-field-as-literal-map."""
mapping_path = tmp_path / "enum_field_as_literal_map.json"
mapping_path.write_text("{invalid json}", encoding="utf-8")
with pytest.raises(SystemExit) as exc_info:
run_main_with_args([
"--input",
str(JSON_SCHEMA_DATA_PATH / "enum_field_as_literal_map.json"),
"--output",
str(tmp_path / "output.py"),
"--enum-field-as-literal-map",
str(mapping_path),
])
assert exc_info.value.code == 2
captured = capsys.readouterr()
assert "Invalid JSON:" in captured.err


def test_main_enum_field_as_literal_map_invalid_encoding_json_file(
tmp_path: Path, capsys: pytest.CaptureFixture[str]
) -> None:
"""Test invalid-encoding JSON file passed to --enum-field-as-literal-map."""
mapping_path = tmp_path / "enum_field_as_literal_map.json"
mapping_path.write_bytes(b"\x80")
with pytest.raises(SystemExit) as exc_info:
run_main_with_args([
"--input",
str(JSON_SCHEMA_DATA_PATH / "enum_field_as_literal_map.json"),
"--output",
str(tmp_path / "output.py"),
"--enum-field-as-literal-map",
str(mapping_path),
])
assert exc_info.value.code == 2
captured = capsys.readouterr()
assert "Unable to read JSON file" in captured.err


def test_main_x_enum_field_as_literal(output_file: Path) -> None:
"""Test x-enum-field-as-literal schema extension for per-field control."""
run_main_and_assert(
Expand Down
Loading