FilesExpand file tree

 main

/

cli.py

Copy path

Blame

More file actions

Blame

More file actions

Latest commit

 

History

History

History

129 lines (103 loc) · 4.09 KB

 main

/

cli.py

Copy path

Top

File metadata and controls

• Code
• Blame

129 lines (103 loc) · 4.09 KB

Raw

Copy raw file

Download raw file

Open symbols panel

Edit and raw actions

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

64

65

66

67

68

69

70

71

72

73

74

75

76

77

78

79

80

81

82

83

84

85

86

87

88

89

90

91

92

93

94

95

96

97

98

99

100

101

102

103

104

105

106

107

108

109

110

111

112

113

114

115

116

117

118

119

120

121

122

123

124

125

126

127

128

129

from __future__ import annotations

import argparse

import json

from typing import List, Literal, Union, Any, Type, TypeVar, Dict

from pydantic import BaseModel

def _get_base_type(annotation: Type[Any]) -> Type[Any]:

if getattr(annotation, "__origin__", None) is Literal:

assert hasattr(annotation, "__args__") and len(annotation.__args__) >= 1 # type: ignore

return type(annotation.__args__[0]) # type: ignore

elif getattr(annotation, "__origin__", None) is Union:

assert hasattr(annotation, "__args__") and len(annotation.__args__) >= 1 # type: ignore

non_optional_args: List[Type[Any]] = [

arg

for arg in annotation.__args__

if arg is not type(None) # type: ignore

]

if non_optional_args:

return _get_base_type(non_optional_args[0])

elif (

getattr(annotation, "__origin__", None) is list

or getattr(annotation, "__origin__", None) is List

):

assert hasattr(annotation, "__args__") and len(annotation.__args__) >= 1 # type: ignore

return _get_base_type(annotation.__args__[0]) # type: ignore

return annotation

def _contains_list_type(annotation: Type[Any] | None) -> bool:

origin = getattr(annotation, "__origin__", None)

if origin is list or origin is List:

return True

elif origin in (Literal, Union):

return any(_contains_list_type(arg) for arg in annotation.__args__) # type: ignore

else:

return False

def _contains_dict_type(annotation: Type[Any] | None) -> bool:

origin = getattr(annotation, "__origin__", None)

if origin is dict or origin is Dict:

return True

elif origin in (Literal, Union):

return any(_contains_dict_type(arg) for arg in annotation.__args__) # type: ignore

else:

return False

def _parse_bool_arg(arg: str | bytes | bool) -> bool:

if isinstance(arg, bytes):

arg = arg.decode("utf-8")

true_values = {"1", "on", "t", "true", "y", "yes"}

false_values = {"0", "off", "f", "false", "n", "no"}

arg_str = str(arg).lower().strip()

if arg_str in true_values:

return True

elif arg_str in false_values:

return False

else:

raise ValueError(f"Invalid boolean argument: {arg}")

def _parse_json_object_arg(arg: str | bytes) -> dict[str, Any]:

if isinstance(arg, bytes):

arg = arg.decode("utf-8")

value = json.loads(arg)

if not isinstance(value, dict):

raise ValueError(f"Invalid JSON object argument: {arg}")

return value

def add_args_from_model(parser: argparse.ArgumentParser, model: Type[BaseModel]):

"""Add arguments from a pydantic model to an argparse parser."""

for name, field in model.model_fields.items():

description = field.description

if field.default and description and not field.is_required():

description += f" (default: {field.default})"

base_type = (

_get_base_type(field.annotation) if field.annotation is not None else str

)

list_type = _contains_list_type(field.annotation)

dict_type = _contains_dict_type(field.annotation)

if dict_type:

parser.add_argument(

f"--{name}",

dest=name,

type=_parse_json_object_arg,

help=description,

)

elif base_type is not bool:

parser.add_argument(

f"--{name}",

dest=name,

nargs="*" if list_type else None,

type=base_type,

help=description,

)

if base_type is bool:

parser.add_argument(

f"--{name}",

dest=name,

type=_parse_bool_arg,

help=f"{description}",

)

T = TypeVar("T", bound=Type[BaseModel])

def parse_model_from_args(model: T, args: argparse.Namespace) -> T:

"""Parse a pydantic model from an argparse namespace."""

return model(

**{

k: v

for k, v in vars(args).items()

if v is not None and k in model.model_fields

}

)