You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
 
 
 
 

140 lines
4.2 KiB

  1. from __future__ import annotations
  2. import os
  3. import warnings
  4. from collections.abc import Iterator, Mapping, MutableMapping
  5. from pathlib import Path
  6. from typing import Any, Callable, TypeVar, overload
  7. class undefined:
  8. pass
  9. class EnvironError(Exception):
  10. pass
  11. class Environ(MutableMapping[str, str]):
  12. def __init__(self, environ: MutableMapping[str, str] = os.environ):
  13. self._environ = environ
  14. self._has_been_read: set[str] = set()
  15. def __getitem__(self, key: str) -> str:
  16. self._has_been_read.add(key)
  17. return self._environ.__getitem__(key)
  18. def __setitem__(self, key: str, value: str) -> None:
  19. if key in self._has_been_read:
  20. raise EnvironError(f"Attempting to set environ['{key}'], but the value has already been read.")
  21. self._environ.__setitem__(key, value)
  22. def __delitem__(self, key: str) -> None:
  23. if key in self._has_been_read:
  24. raise EnvironError(f"Attempting to delete environ['{key}'], but the value has already been read.")
  25. self._environ.__delitem__(key)
  26. def __iter__(self) -> Iterator[str]:
  27. return iter(self._environ)
  28. def __len__(self) -> int:
  29. return len(self._environ)
  30. environ = Environ()
  31. T = TypeVar("T")
  32. class Config:
  33. def __init__(
  34. self,
  35. env_file: str | Path | None = None,
  36. environ: Mapping[str, str] = environ,
  37. env_prefix: str = "",
  38. ) -> None:
  39. self.environ = environ
  40. self.env_prefix = env_prefix
  41. self.file_values: dict[str, str] = {}
  42. if env_file is not None:
  43. if not os.path.isfile(env_file):
  44. warnings.warn(f"Config file '{env_file}' not found.")
  45. else:
  46. self.file_values = self._read_file(env_file)
  47. @overload
  48. def __call__(self, key: str, *, default: None) -> str | None: ...
  49. @overload
  50. def __call__(self, key: str, cast: type[T], default: T = ...) -> T: ...
  51. @overload
  52. def __call__(self, key: str, cast: type[str] = ..., default: str = ...) -> str: ...
  53. @overload
  54. def __call__(
  55. self,
  56. key: str,
  57. cast: Callable[[Any], T] = ...,
  58. default: Any = ...,
  59. ) -> T: ...
  60. @overload
  61. def __call__(self, key: str, cast: type[str] = ..., default: T = ...) -> T | str: ...
  62. def __call__(
  63. self,
  64. key: str,
  65. cast: Callable[[Any], Any] | None = None,
  66. default: Any = undefined,
  67. ) -> Any:
  68. return self.get(key, cast, default)
  69. def get(
  70. self,
  71. key: str,
  72. cast: Callable[[Any], Any] | None = None,
  73. default: Any = undefined,
  74. ) -> Any:
  75. key = self.env_prefix + key
  76. if key in self.environ:
  77. value = self.environ[key]
  78. return self._perform_cast(key, value, cast)
  79. if key in self.file_values:
  80. value = self.file_values[key]
  81. return self._perform_cast(key, value, cast)
  82. if default is not undefined:
  83. return self._perform_cast(key, default, cast)
  84. raise KeyError(f"Config '{key}' is missing, and has no default.")
  85. def _read_file(self, file_name: str | Path) -> dict[str, str]:
  86. file_values: dict[str, str] = {}
  87. with open(file_name) as input_file:
  88. for line in input_file.readlines():
  89. line = line.strip()
  90. if "=" in line and not line.startswith("#"):
  91. key, value = line.split("=", 1)
  92. key = key.strip()
  93. value = value.strip().strip("\"'")
  94. file_values[key] = value
  95. return file_values
  96. def _perform_cast(
  97. self,
  98. key: str,
  99. value: Any,
  100. cast: Callable[[Any], Any] | None = None,
  101. ) -> Any:
  102. if cast is None or value is None:
  103. return value
  104. elif cast is bool and isinstance(value, str):
  105. mapping = {"true": True, "1": True, "false": False, "0": False}
  106. value = value.lower()
  107. if value not in mapping:
  108. raise ValueError(f"Config '{key}' has value '{value}'. Not a valid bool.")
  109. return mapping[value]
  110. try:
  111. return cast(value)
  112. except (TypeError, ValueError):
  113. raise ValueError(f"Config '{key}' has value '{value}'. Not a valid {cast.__name__}.")