"""Pydantic models for argparse arguments and argument groups.""" from typing import ( Any, Literal, ) from pydantic import BaseModel, model_validator SEP = "__" class ArgparseArgumentModel(BaseModel): """Pydantic model for an argparse argument.""" name: str type: Any dest: str default: Any required: bool action: Literal["store_true", "store"] help: str | None nargs: Literal["+"] | None choices: tuple | None @model_validator(mode="after") # ty: ignore[invalid-argument-type] @classmethod def validate_action(cls, values: "ArgparseArgumentModel"): """Validate the action based on the type.""" if values.type is bool and values.action != "store_true": raise ValueError('If type is bool, action must be "store_true"') return values @model_validator(mode="after") # ty: ignore[invalid-argument-type] @classmethod def remove_props_on_store_true(cls, values: "ArgparseArgumentModel"): """Remove type, nargs, and choices if action is store_true.""" if values.action == "store_true": values.type = None values.nargs = None values.choices = None return values def model_dump(self, **kwargs): """Override the model_dump method to remove empty choices.""" res = super().model_dump(**kwargs) if "choices" in res and not res["choices"]: del res["choices"] return res class ArgparseArgumentGroupModel(BaseModel): """Pydantic model for a custom argument group.""" name: str arguments: list[ArgparseArgumentModel]