1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
|
"""Pydantic models for argparse arguments and argument groups."""
from typing import (
Any,
List,
Literal,
Optional,
Tuple,
)
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: Optional[str]
nargs: Optional[Literal["+"]]
choices: Optional[Tuple]
@model_validator(mode="after") # type: ignore
@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") # type: ignore
@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
# override
def model_dump(self, **kwargs):
"""Override the model_dump method to remove empty choices."""
res = super().model_dump(**kwargs)
# Check if choices is present and if it's an empty tuple remove it
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]
|
