|
9 | 9 | from collections.abc import Callable |
10 | 10 |
|
11 | 11 | BASE_DIR: Final[Path] = Path(__file__).parent.parent |
12 | | -TRUE_VALUES: Final[frozenset[str]] = frozenset({"True", "true", "1", "yes", "Y", "T"}) |
| 12 | +TRUE_VALUES: Final[frozenset[str]] = frozenset({"True", "true", "1", "yes", "YES", "Y", "y", "T", "t"}) |
13 | 13 |
|
14 | 14 | T = TypeVar("T") |
15 | | -ParseTypes = bool | int | str | list[str] |
| 15 | +ParseTypes = bool | int | str | list[str] | Path | list[Path] |
| 16 | + |
| 17 | + |
| 18 | +class UnsetType: |
| 19 | + """Placeholder for an Unset type. |
| 20 | +
|
| 21 | + This helps differentiate None from the default |
| 22 | + """ |
| 23 | + |
| 24 | + |
| 25 | +_UNSET = UnsetType() |
| 26 | + |
| 27 | + |
| 28 | +@overload |
| 29 | +def get_env(key: str, default: bool, type_hint: UnsetType = _UNSET) -> Callable[[], bool]: ... |
| 30 | + |
| 31 | + |
| 32 | +@overload |
| 33 | +def get_env(key: str, default: int, type_hint: UnsetType = _UNSET) -> Callable[[], int]: ... |
| 34 | + |
| 35 | + |
| 36 | +@overload |
| 37 | +def get_env(key: str, default: str, type_hint: UnsetType = _UNSET) -> Callable[[], str]: ... |
| 38 | + |
| 39 | + |
| 40 | +@overload |
| 41 | +def get_env(key: str, default: Path, type_hint: UnsetType = _UNSET) -> Callable[[], Path]: ... |
16 | 42 |
|
17 | 43 |
|
18 | 44 | @overload |
19 | | -def get_env(key: str, default: bool) -> Callable[[], bool]: ... |
| 45 | +def get_env(key: str, default: list[Path], type_hint: UnsetType = _UNSET) -> Callable[[], list[Path]]: ... |
20 | 46 |
|
21 | 47 |
|
22 | 48 | @overload |
23 | | -def get_env(key: str, default: int) -> Callable[[], int]: ... |
| 49 | +def get_env(key: str, default: list[str], type_hint: UnsetType = _UNSET) -> Callable[[], list[str]]: ... |
24 | 50 |
|
25 | 51 |
|
26 | 52 | @overload |
27 | | -def get_env(key: str, default: str) -> Callable[[], str]: ... |
| 53 | +def get_env(key: str, default: None, type_hint: UnsetType = _UNSET) -> Callable[[], None]: ... |
28 | 54 |
|
29 | 55 |
|
30 | 56 | @overload |
31 | | -def get_env(key: str, default: list[str]) -> Callable[[], list[str]]: ... |
| 57 | +def get_env(key: str, default: ParseTypes | None, type_hint: type[T]) -> Callable[[], T]: ... |
32 | 58 |
|
33 | 59 |
|
34 | | -def get_env(key: str, default: ParseTypes) -> Callable[[], ParseTypes]: |
35 | | - return lambda: get_config_val(key, default) |
| 60 | +def get_env( |
| 61 | + key: str, default: ParseTypes | None, type_hint: type[T] | UnsetType = _UNSET |
| 62 | +) -> Callable[[], ParseTypes | T | None]: |
| 63 | + return lambda: get_config_val(key=key, default=default, type_hint=type_hint) |
36 | 64 |
|
37 | 65 |
|
38 | 66 | @overload |
39 | | -def get_config_val(key: str, default: bool) -> bool: ... |
| 67 | +def get_config_val(key: str, default: bool, type_hint: UnsetType = _UNSET) -> bool: ... |
40 | 68 |
|
41 | 69 |
|
42 | 70 | @overload |
43 | | -def get_config_val(key: str, default: int) -> int: ... |
| 71 | +def get_config_val(key: str, default: int, type_hint: UnsetType = _UNSET) -> int: ... |
44 | 72 |
|
45 | 73 |
|
46 | 74 | @overload |
47 | | -def get_config_val(key: str, default: str) -> str: ... |
| 75 | +def get_config_val(key: str, default: str, type_hint: UnsetType = _UNSET) -> str: ... |
48 | 76 |
|
49 | 77 |
|
50 | 78 | @overload |
51 | | -def get_config_val(key: str, default: list[str]) -> list[str]: ... |
| 79 | +def get_config_val(key: str, default: Path, type_hint: UnsetType = _UNSET) -> Path: ... |
52 | 80 |
|
53 | 81 |
|
54 | | -def get_config_val( |
55 | | - key: str, |
56 | | - default: ParseTypes, |
57 | | -) -> ParseTypes: |
| 82 | +@overload |
| 83 | +def get_config_val(key: str, default: list[Path], type_hint: UnsetType = _UNSET) -> list[Path]: ... |
| 84 | + |
| 85 | + |
| 86 | +@overload |
| 87 | +def get_config_val(key: str, default: list[str], type_hint: UnsetType = _UNSET) -> list[str]: ... |
| 88 | + |
| 89 | + |
| 90 | +@overload |
| 91 | +def get_config_val(key: str, default: None, type_hint: UnsetType = _UNSET) -> None: ... |
| 92 | + |
| 93 | + |
| 94 | +@overload |
| 95 | +def get_config_val(key: str, default: ParseTypes | None, type_hint: type[T]) -> T: ... |
| 96 | + |
| 97 | + |
| 98 | +def get_config_val( # noqa: C901, PLR0912, PLR0911 |
| 99 | + key: str, default: ParseTypes | None, type_hint: type[T] | UnsetType = _UNSET |
| 100 | +) -> ParseTypes | T | None: |
58 | 101 | """Parse environment variables. |
59 | 102 |
|
60 | 103 | Args: |
61 | 104 | key: Environment variable key |
62 | 105 | default: Default value if key not found in environment |
| 106 | + type_hint: Optional type hint to use instead of inferring from `default` |
| 107 | +
|
| 108 | + Raises: |
| 109 | + ValueError: Raised when the configuration value cannot be parsed. |
63 | 110 |
|
64 | 111 | Returns: |
65 | 112 | Parsed value of the specified type |
66 | 113 | """ |
67 | | - value = os.getenv(key) |
68 | | - if value is None: |
| 114 | + str_value = os.getenv(key) |
| 115 | + if str_value is None: |
| 116 | + if type_hint != _UNSET: |
| 117 | + return cast("T", default) |
69 | 118 | return default |
| 119 | + value: str = str_value |
70 | 120 | if type(default) is bool: |
71 | | - return value in TRUE_VALUES |
| 121 | + bool_value = value in TRUE_VALUES |
| 122 | + if type_hint != _UNSET: |
| 123 | + return cast("T", bool_value) |
| 124 | + return bool_value |
72 | 125 | if type(default) is int: |
73 | | - return int(value) |
74 | | - if type(default) is list: |
| 126 | + int_value = int(value) |
| 127 | + if type_hint != _UNSET: |
| 128 | + return cast("T", int_value) |
| 129 | + return int_value |
| 130 | + if type(default) is Path: |
| 131 | + path_value = Path(value) |
| 132 | + if type_hint != _UNSET: |
| 133 | + return cast("T", path_value) |
| 134 | + return path_value |
| 135 | + if type(default) is list[Path]: |
| 136 | + if value.startswith("[") and value.endswith("]"): |
| 137 | + try: |
| 138 | + path_list = [Path(s) for s in json.loads(value)] |
| 139 | + if type_hint != _UNSET: |
| 140 | + return cast("T", path_list) |
| 141 | + except (SyntaxError, ValueError) as e: |
| 142 | + msg = f"{key} is not a valid list representation." |
| 143 | + raise ValueError(msg) from e |
| 144 | + else: |
| 145 | + return value |
| 146 | + path_list = [Path(host.strip()) for host in value.split(",")] |
| 147 | + if type_hint != _UNSET: |
| 148 | + return cast("T", path_list) |
| 149 | + return path_list |
| 150 | + if type(default) is list[str]: |
75 | 151 | if value.startswith("[") and value.endswith("]"): |
76 | 152 | try: |
77 | | - return cast("list[str]", json.loads(value)) |
| 153 | + str_list = cast("list[str]", json.loads(value)) |
| 154 | + if type_hint != _UNSET: |
| 155 | + return cast("T", str_list) |
78 | 156 | except (SyntaxError, ValueError) as e: |
79 | 157 | msg = f"{key} is not a valid list representation." |
80 | 158 | raise ValueError(msg) from e |
81 | | - return [host.strip() for host in value.split(",")] |
| 159 | + else: |
| 160 | + return value |
| 161 | + str_list = [host.strip() for host in value.split(",")] |
| 162 | + if type_hint != _UNSET: |
| 163 | + return cast("T", str_list) |
| 164 | + return str_list |
| 165 | + if type_hint != _UNSET: |
| 166 | + return cast("T", value) |
82 | 167 | return value |
0 commit comments