config.py 4.5 KB

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