datamodel-code-generator 中整数枚举值处理问题解析

问题背景

在开发过程中，当使用 datamodel-code-generator 工具从 Swagger/OpenAPI 规范生成 Pydantic 模型时，遇到了一个关于整数枚举值处理的错误。具体场景是当 Swagger 规范中定义了一个带有整数枚举值的属性时，代码生成过程会抛出异常。

问题重现

考虑以下 Swagger 规范片段：

properties:
  block_size:
    default: 4096
    description: "The size of blocks in the namespace, in bytes."
    enum:
    - 512
    - 4096
    type: integer

当使用 datamodel-code-generator 工具生成 Pydantic 模型时，会抛出以下错误：

AttributeError: 'int' object has no attribute 'strip'

问题分析

这个错误的根本原因在于枚举处理逻辑中的一个假设错误。在 datamodel_code_generator/model/enum.py 文件的第 94 行，代码假设所有枚举成员的值都是字符串类型，并尝试对这些值调用 strip() 方法：

field_default = (field.default or "").strip("'\"")

然而，当枚举值是整数类型时（如示例中的 512 和 4096），这个假设就不成立了，因为整数类型没有 strip() 方法。

解决方案

修复方案相对简单：需要先检查字段默认值的类型，如果是字符串才调用 strip() 方法，否则直接转换为字符串：

if isinstance(field.default, str):
    field_default = field.default.strip("'\"")
else:
    field_default = str(field.default)

这个修改确保了无论枚举值是字符串还是整数类型，都能正确处理。

预期输出

修复后，工具应该能够正确生成类似以下的 Pydantic 模型代码：

class BlockSize(Enum):
    integer_512 = 512
    integer_4096 = 4096

class Space3(BaseModel):
    block_size: BlockSize | None = Field(
        BlockSize.integer_4096,
        description="The size of blocks in the namespace, in bytes.",
    )

技术影响

这个修复不会引入任何负面副作用，因为它：

1. 保持了原有字符串处理的逻辑不变
2. 只是增加了对非字符串类型的处理
3. 最终结果仍然是字符串比较，与原有逻辑一致

最佳实践建议

在使用 datamodel-code-generator 处理枚举值时，建议：

1. 明确指定枚举值的类型
2. 对于数值型枚举，确保在 Swagger 规范中正确设置 type 字段
3. 测试生成的代码以确保枚举值被正确识别和处理

总结

这个问题展示了类型假设在代码生成工具中的重要性。工具开发者需要考虑到各种可能的输入类型，特别是在处理来自不同来源的规范时。通过添加类型检查，我们能够使工具更加健壮，能够处理更广泛的使用场景。