config.py 4.3 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140
  1. from __future__ import annotations
  2. import os
  3. import warnings
  4. from collections.abc import Callable, Iterator, Mapping, MutableMapping
  5. from pathlib import Path
  6. from typing import Any, 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. encoding: str = "utf-8",
  39. ) -> None:
  40. self.environ = environ
  41. self.env_prefix = env_prefix
  42. self.file_values: dict[str, str] = {}
  43. if env_file is not None:
  44. if not os.path.isfile(env_file):
  45. warnings.warn(f"Config file '{env_file}' not found.")
  46. else:
  47. self.file_values = self._read_file(env_file, encoding)
  48. @overload
  49. def __call__(self, key: str, *, default: None) -> str | None: ...
  50. @overload
  51. def __call__(self, key: str, cast: type[T], default: T = ...) -> T: ...
  52. @overload
  53. def __call__(self, key: str, cast: type[str] = ..., default: str = ...) -> str: ...
  54. @overload
  55. def __call__(
  56. self,
  57. key: str,
  58. cast: Callable[[Any], T] = ...,
  59. default: Any = ...,
  60. ) -> T: ...
  61. @overload
  62. def __call__(self, key: str, cast: type[str] = ..., default: T = ...) -> T | str: ...
  63. def __call__(
  64. self,
  65. key: str,
  66. cast: Callable[[Any], Any] | None = None,
  67. default: Any = undefined,
  68. ) -> Any:
  69. return self.get(key, cast, default)
  70. def get(
  71. self,
  72. key: str,
  73. cast: Callable[[Any], Any] | None = None,
  74. default: Any = undefined,
  75. ) -> Any:
  76. key = self.env_prefix + key
  77. if key in self.environ:
  78. value = self.environ[key]
  79. return self._perform_cast(key, value, cast)
  80. if key in self.file_values:
  81. value = self.file_values[key]
  82. return self._perform_cast(key, value, cast)
  83. if default is not undefined:
  84. return self._perform_cast(key, default, cast)
  85. raise KeyError(f"Config '{key}' is missing, and has no default.")
  86. def _read_file(self, file_name: str | Path, encoding: str) -> dict[str, str]:
  87. file_values: dict[str, str] = {}
  88. with open(file_name, encoding=encoding) as input_file:
  89. for line in input_file.readlines():
  90. line = line.strip()
  91. if "=" in line and not line.startswith("#"):
  92. key, value = line.split("=", 1)
  93. key = key.strip()
  94. value = value.strip().strip("\"'")
  95. file_values[key] = value
  96. return file_values
  97. def _perform_cast(
  98. self,
  99. key: str,
  100. value: Any,
  101. cast: Callable[[Any], Any] | None = None,
  102. ) -> Any:
  103. if cast is None or value is None:
  104. return value
  105. elif cast is bool and isinstance(value, str):
  106. mapping = {"true": True, "1": True, "false": False, "0": False}
  107. value = value.lower()
  108. if value not in mapping:
  109. raise ValueError(f"Config '{key}' has value '{value}'. Not a valid bool.")
  110. return mapping[value]
  111. try:
  112. return cast(value)
  113. except (TypeError, ValueError):
  114. raise ValueError(f"Config '{key}' has value '{value}'. Not a valid {cast.__name__}.")