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.
 
 
 
 

130 lines
4.1 KiB

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