diff --git a/src/datamodel_code_generator/format.py b/src/datamodel_code_generator/format.py index 7cdda90ab..4338af841 100644 --- a/src/datamodel_code_generator/format.py +++ b/src/datamodel_code_generator/format.py @@ -74,6 +74,12 @@ class PythonVersion(Enum): PY_313 = "3.13" PY_314 = "3.14" + @cached_property + def version_key(self) -> tuple[int, int]: + """Return (major, minor) tuple for version comparison.""" + major, minor = self.value.split(".") + return int(major), int(minor) + @cached_property def _is_py_310_or_later(self) -> bool: # pragma: no cover return True # 3.10+ always true since minimum is PY_310 diff --git a/src/datamodel_code_generator/parser/base.py b/src/datamodel_code_generator/parser/base.py index 0436fe89d..d5cad044e 100644 --- a/src/datamodel_code_generator/parser/base.py +++ b/src/datamodel_code_generator/parser/base.py @@ -126,17 +126,12 @@ def __ge__(self, value: Any, /) -> bool: ... # noqa: D105 } -def _python_version_key(python_version: PythonVersion) -> tuple[int, int]: - major, minor = python_version.value.split(".") - return int(major), int(minor) - - def _get_builtin_names_for_target(target_python_version: PythonVersion) -> frozenset[str]: builtin_names = set(_BUILTIN_NAMES) - target_key = _python_version_key(target_python_version) + target_key = target_python_version.version_key for introduced_version, names in _BUILTIN_NAMES_INTRODUCED_IN.items(): - if target_key >= _python_version_key(introduced_version): + if target_key >= introduced_version.version_key: builtin_names.update(names) else: builtin_names.difference_update(names) @@ -1061,6 +1056,8 @@ def __init__( # noqa: PLR0912, PLR0915 remove_special_field_name_prefix=config.remove_special_field_name_prefix, capitalise_enum_members=config.capitalise_enum_members, no_alias=config.no_alias, + use_subclass_enum=config.use_subclass_enum, + target_python_version=config.target_python_version, parent_scoped_naming=config.parent_scoped_naming, treat_dot_as_module=config.treat_dot_as_module, naming_strategy=config.naming_strategy, diff --git a/src/datamodel_code_generator/reference.py b/src/datamodel_code_generator/reference.py index 9ab8afdb5..189a59efc 100644 --- a/src/datamodel_code_generator/reference.py +++ b/src/datamodel_code_generator/reference.py @@ -36,6 +36,7 @@ from datamodel_code_generator import Error, NamingStrategy from datamodel_code_generator.enums import ClassNameAffixScope +from datamodel_code_generator.format import PythonVersion from datamodel_code_generator.util import ConfigDict, camel_to_snake, is_pydantic_v2, model_validator if TYPE_CHECKING: @@ -221,6 +222,27 @@ def context_variable(setter: Callable[[T], None], current_value: T, new_value: T setter(previous_value) +_BUILTIN_TYPE_ATTRIBUTES: frozenset[str] = frozenset( + name for cls in (str, int, float, bytes) for name in dir(cls) if not name.startswith("_") +) + +_BUILTIN_TYPE_ATTRIBUTES_INTRODUCED_IN: dict[PythonVersion, frozenset[str]] = { + PythonVersion.PY_314: frozenset({"from_number"}), +} + + +def _get_builtin_type_attributes_for_target(target: PythonVersion) -> frozenset[str]: + """Get builtin type attributes adjusted for the target Python version.""" + attrs = set(_BUILTIN_TYPE_ATTRIBUTES) + target_key = target.version_key + for ver, names in _BUILTIN_TYPE_ATTRIBUTES_INTRODUCED_IN.items(): + if target_key >= ver.version_key: + attrs.update(names) + else: + attrs.difference_update(names) + return frozenset(attrs) + + class FieldNameResolver: """Converts schema field names to valid Python identifiers.""" @@ -234,6 +256,8 @@ def __init__( # noqa: PLR0913, PLR0917 remove_special_field_name_prefix: bool = False, # noqa: FBT001, FBT002 capitalise_enum_members: bool = False, # noqa: FBT001, FBT002 no_alias: bool = False, # noqa: FBT001, FBT002 + use_subclass_enum: bool = False, # noqa: FBT001, FBT002 + target_python_version: PythonVersion | None = None, ) -> None: """Initialize field name resolver with transformation options.""" self.aliases: Mapping[str, str | list[str]] = {} if aliases is None else {**aliases} @@ -246,9 +270,10 @@ def __init__( # noqa: PLR0913, PLR0917 self.remove_special_field_name_prefix: bool = remove_special_field_name_prefix self.capitalise_enum_members: bool = capitalise_enum_members self.no_alias = no_alias + self.use_subclass_enum: bool = use_subclass_enum + self.target_python_version = target_python_version - @classmethod - def _validate_field_name(cls, field_name: str) -> bool: # noqa: ARG003 + def _validate_field_name(self, field_name: str) -> bool: # noqa: ARG002, PLR6301 """Check if a field name is valid. Subclasses may override.""" return True @@ -283,7 +308,8 @@ def get_valid_name( # noqa: PLR0912 if self.capitalise_enum_members or (self.snake_case_field and not ignore_snake_case_field): name = camel_to_snake(name) count = 1 - if iskeyword(name) or not self._validate_field_name(name): + validated_name = name.upper() if self.capitalise_enum_members else name + if iskeyword(validated_name) or not self._validate_field_name(validated_name): name += "_" if upper_camel: new_name = snake_to_upper_camel(name) @@ -357,100 +383,28 @@ def get_valid_field_name_and_alias( class PydanticFieldNameResolver(FieldNameResolver): """Field name resolver that avoids Pydantic reserved names.""" - @classmethod - def _validate_field_name(cls, field_name: str) -> bool: + def _validate_field_name(self, field_name: str) -> bool: # noqa: PLR6301 """Check field name doesn't conflict with BaseModel attributes.""" # TODO: Support Pydantic V2 return not hasattr(BaseModel, field_name) class EnumFieldNameResolver(FieldNameResolver): - """Field name resolver for enum members with special handling for reserved names. - - When using --use-subclass-enum, enums inherit from types like str or int. - Member names that conflict with methods of these types cause type checker errors. - This class detects and handles such conflicts by adding underscore suffixes. - - The _BUILTIN_TYPE_ATTRIBUTES set is intentionally static (not using hasattr) - to avoid runtime Python version differences affecting code generation. - Based on Python 3.8-3.14 method names (union of all versions for safety). - Note: 'mro' is handled explicitly in get_valid_name for backward compatibility. - """ - - _BUILTIN_TYPE_ATTRIBUTES: ClassVar[frozenset[str]] = frozenset({ - "as_integer_ratio", - "bit_count", - "bit_length", - "capitalize", - "casefold", - "center", - "conjugate", - "count", - "decode", - "denominator", - "encode", - "endswith", - "expandtabs", - "find", - "format", - "format_map", - "from_bytes", - "from_number", - "fromhex", - "hex", - "imag", - "index", - "isalnum", - "isalpha", - "isascii", - "isdecimal", - "isdigit", - "isidentifier", - "islower", - "isnumeric", - "isprintable", - "isspace", - "istitle", - "isupper", - "is_integer", - "join", - "ljust", - "lower", - "lstrip", - "maketrans", - "numerator", - "partition", - "real", - "removeprefix", - "removesuffix", - "replace", - "rfind", - "rindex", - "rjust", - "rpartition", - "rsplit", - "rstrip", - "split", - "splitlines", - "startswith", - "strip", - "swapcase", - "title", - "to_bytes", - "translate", - "upper", - "zfill", - }) - - @classmethod - def _validate_field_name(cls, field_name: str) -> bool: - """Check field name doesn't conflict with subclass enum base type attributes. + """Field name resolver for enum members with special handling for reserved names.""" + + def __init__(self, **kwargs: Any) -> None: + """Initialize with version-aware builtin type attributes.""" + super().__init__(**kwargs) + target = self.target_python_version + self._builtin_type_attributes = ( + _get_builtin_type_attributes_for_target(target) if target is not None else _BUILTIN_TYPE_ATTRIBUTES + ) - When using --use-subclass-enum, enums inherit from types like str or int. - Member names that conflict with methods of these types (e.g., 'count' for str) - cause type checker errors. This method detects such conflicts. - """ - return field_name not in cls._BUILTIN_TYPE_ATTRIBUTES + def _validate_field_name(self, field_name: str) -> bool: + """Check field name doesn't conflict with subclass enum base type attributes.""" + if not self.use_subclass_enum: + return True + return field_name not in self._builtin_type_attributes def get_valid_name( self, @@ -538,6 +492,8 @@ def __init__( # noqa: PLR0913, PLR0917 remove_special_field_name_prefix: bool = False, # noqa: FBT001, FBT002 capitalise_enum_members: bool = False, # noqa: FBT001, FBT002 no_alias: bool = False, # noqa: FBT001, FBT002 + use_subclass_enum: bool = False, # noqa: FBT001, FBT002 + target_python_version: PythonVersion | None = None, remove_suffix_number: bool = False, # noqa: FBT001, FBT002 parent_scoped_naming: bool = False, # noqa: FBT001, FBT002 treat_dot_as_module: bool | None = None, # noqa: FBT001 @@ -575,6 +531,8 @@ def __init__( # noqa: PLR0913, PLR0917 remove_special_field_name_prefix=remove_special_field_name_prefix, capitalise_enum_members=capitalise_enum_members if k == ModelType.ENUM else False, no_alias=no_alias, + use_subclass_enum=use_subclass_enum if k == ModelType.ENUM else False, + target_python_version=target_python_version if k == ModelType.ENUM else None, ) for k, v in merged_field_name_resolver_classes.items() } diff --git a/tests/data/expected/main_kr/capitalise_enum_members_and_use_subclass_enum_builtin_conflict.py b/tests/data/expected/main_kr/capitalise_enum_members_and_use_subclass_enum_builtin_conflict.py new file mode 100644 index 000000000..2a2c67cc3 --- /dev/null +++ b/tests/data/expected/main_kr/capitalise_enum_members_and_use_subclass_enum_builtin_conflict.py @@ -0,0 +1,11 @@ +# generated by datamodel-codegen: +# filename: enum_builtin_conflict_two.json + +from __future__ import annotations + +from enum import Enum + + +class Model(str, Enum): + REPLACE = 'replace' + COUNT = 'count' diff --git a/tests/data/expected/main_kr/capitalise_enum_members_builtin_conflict.py b/tests/data/expected/main_kr/capitalise_enum_members_builtin_conflict.py new file mode 100644 index 000000000..3eb6df657 --- /dev/null +++ b/tests/data/expected/main_kr/capitalise_enum_members_builtin_conflict.py @@ -0,0 +1,12 @@ +# generated by datamodel-codegen: +# filename: enum_builtin_conflict.json + +from __future__ import annotations + +from enum import Enum + + +class Model(Enum): + REPLACE = 'replace' + COUNT = 'count' + INDEX = 'index' diff --git a/tests/data/expected/main_kr/no_subclass_enum_no_capitalise_builtin_names.py b/tests/data/expected/main_kr/no_subclass_enum_no_capitalise_builtin_names.py new file mode 100644 index 000000000..b411dc4e5 --- /dev/null +++ b/tests/data/expected/main_kr/no_subclass_enum_no_capitalise_builtin_names.py @@ -0,0 +1,11 @@ +# generated by datamodel-codegen: +# filename: enum_builtin_conflict_two.json + +from __future__ import annotations + +from enum import Enum + + +class Model(Enum): + replace = 'replace' + count = 'count' diff --git a/tests/data/expected/main_kr/use_subclass_enum_builtin_conflict_no_capitalise.py b/tests/data/expected/main_kr/use_subclass_enum_builtin_conflict_no_capitalise.py new file mode 100644 index 000000000..4b610cddb --- /dev/null +++ b/tests/data/expected/main_kr/use_subclass_enum_builtin_conflict_no_capitalise.py @@ -0,0 +1,11 @@ +# generated by datamodel-codegen: +# filename: enum_builtin_conflict_two.json + +from __future__ import annotations + +from enum import Enum + + +class Model(str, Enum): + replace_ = 'replace' + count_ = 'count' diff --git a/tests/data/jsonschema/enum_builtin_conflict.json b/tests/data/jsonschema/enum_builtin_conflict.json new file mode 100644 index 000000000..c684c87fd --- /dev/null +++ b/tests/data/jsonschema/enum_builtin_conflict.json @@ -0,0 +1,4 @@ +{ + "type": "string", + "enum": ["replace", "count", "index"] +} diff --git a/tests/data/jsonschema/enum_builtin_conflict_two.json b/tests/data/jsonschema/enum_builtin_conflict_two.json new file mode 100644 index 000000000..4f8fb044d --- /dev/null +++ b/tests/data/jsonschema/enum_builtin_conflict_two.json @@ -0,0 +1,4 @@ +{ + "type": "string", + "enum": ["replace", "count"] +} diff --git a/tests/test_main_kr.py b/tests/test_main_kr.py index d1923ccfa..b627301e4 100644 --- a/tests/test_main_kr.py +++ b/tests/test_main_kr.py @@ -586,6 +586,70 @@ class EnumSystems(str, Enum): ) +def test_capitalise_enum_members_builtin_conflict(output_file: Path) -> None: + """Test capitalise-enum-members does not add underscore to builtin names (#2970).""" + run_main_and_assert( + input_path=JSON_SCHEMA_DATA_PATH / "enum_builtin_conflict.json", + output_path=output_file, + assert_func=assert_file_content, + input_file_type="jsonschema", + extra_args=[ + "--output-model-type", + "pydantic_v2.BaseModel", + "--disable-timestamp", + "--capitalise-enum-members", + ], + ) + + +def test_capitalise_enum_members_and_use_subclass_enum_builtin_conflict(output_file: Path) -> None: + """Test capitalise-enum-members + use-subclass-enum with builtin names (#2970).""" + run_main_and_assert( + input_path=JSON_SCHEMA_DATA_PATH / "enum_builtin_conflict_two.json", + output_path=output_file, + assert_func=assert_file_content, + input_file_type="jsonschema", + extra_args=[ + "--output-model-type", + "pydantic_v2.BaseModel", + "--disable-timestamp", + "--capitalise-enum-members", + "--use-subclass-enum", + ], + ) + + +def test_use_subclass_enum_builtin_conflict_no_capitalise(output_file: Path) -> None: + """Test use-subclass-enum without capitalise adds underscore for builtin conflicts.""" + run_main_and_assert( + input_path=JSON_SCHEMA_DATA_PATH / "enum_builtin_conflict_two.json", + output_path=output_file, + assert_func=assert_file_content, + input_file_type="jsonschema", + extra_args=[ + "--output-model-type", + "pydantic_v2.BaseModel", + "--disable-timestamp", + "--use-subclass-enum", + ], + ) + + +def test_no_subclass_enum_no_capitalise_builtin_names(output_file: Path) -> None: + """Test default behavior with builtin names has no underscore suffix.""" + run_main_and_assert( + input_path=JSON_SCHEMA_DATA_PATH / "enum_builtin_conflict_two.json", + output_path=output_file, + assert_func=assert_file_content, + input_file_type="jsonschema", + extra_args=[ + "--output-model-type", + "pydantic_v2.BaseModel", + "--disable-timestamp", + ], + ) + + EXPECTED_GENERATE_PYPROJECT_CONFIG_PATH = EXPECTED_MAIN_KR_PATH / "generate_pyproject_config"