main.py 12 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400
  1. import io
  2. import logging
  3. import os
  4. import pathlib
  5. import shutil
  6. import sys
  7. import tempfile
  8. from collections import OrderedDict
  9. from contextlib import contextmanager
  10. from typing import IO, Dict, Iterable, Iterator, Mapping, Optional, Tuple, Union
  11. from .parser import Binding, parse_stream
  12. from .variables import parse_variables
  13. # A type alias for a string path to be used for the paths in this file.
  14. # These paths may flow to `open()` and `shutil.move()`; `shutil.move()`
  15. # only accepts string paths, not byte paths or file descriptors. See
  16. # https://github.com/python/typeshed/pull/6832.
  17. StrPath = Union[str, "os.PathLike[str]"]
  18. logger = logging.getLogger(__name__)
  19. def with_warn_for_invalid_lines(mappings: Iterator[Binding]) -> Iterator[Binding]:
  20. for mapping in mappings:
  21. if mapping.error:
  22. logger.warning(
  23. "python-dotenv could not parse statement starting at line %s",
  24. mapping.original.line,
  25. )
  26. yield mapping
  27. class DotEnv:
  28. def __init__(
  29. self,
  30. dotenv_path: Optional[StrPath],
  31. stream: Optional[IO[str]] = None,
  32. verbose: bool = False,
  33. encoding: Optional[str] = None,
  34. interpolate: bool = True,
  35. override: bool = True,
  36. ) -> None:
  37. self.dotenv_path: Optional[StrPath] = dotenv_path
  38. self.stream: Optional[IO[str]] = stream
  39. self._dict: Optional[Dict[str, Optional[str]]] = None
  40. self.verbose: bool = verbose
  41. self.encoding: Optional[str] = encoding
  42. self.interpolate: bool = interpolate
  43. self.override: bool = override
  44. @contextmanager
  45. def _get_stream(self) -> Iterator[IO[str]]:
  46. if self.dotenv_path and os.path.isfile(self.dotenv_path):
  47. with open(self.dotenv_path, encoding=self.encoding) as stream:
  48. yield stream
  49. elif self.stream is not None:
  50. yield self.stream
  51. else:
  52. if self.verbose:
  53. logger.info(
  54. "python-dotenv could not find configuration file %s.",
  55. self.dotenv_path or ".env",
  56. )
  57. yield io.StringIO("")
  58. def dict(self) -> Dict[str, Optional[str]]:
  59. """Return dotenv as dict"""
  60. if self._dict:
  61. return self._dict
  62. raw_values = self.parse()
  63. if self.interpolate:
  64. self._dict = OrderedDict(
  65. resolve_variables(raw_values, override=self.override)
  66. )
  67. else:
  68. self._dict = OrderedDict(raw_values)
  69. return self._dict
  70. def parse(self) -> Iterator[Tuple[str, Optional[str]]]:
  71. with self._get_stream() as stream:
  72. for mapping in with_warn_for_invalid_lines(parse_stream(stream)):
  73. if mapping.key is not None:
  74. yield mapping.key, mapping.value
  75. def set_as_environment_variables(self) -> bool:
  76. """
  77. Load the current dotenv as system environment variable.
  78. """
  79. if not self.dict():
  80. return False
  81. for k, v in self.dict().items():
  82. if k in os.environ and not self.override:
  83. continue
  84. if v is not None:
  85. os.environ[k] = v
  86. return True
  87. def get(self, key: str) -> Optional[str]:
  88. """ """
  89. data = self.dict()
  90. if key in data:
  91. return data[key]
  92. if self.verbose:
  93. logger.warning("Key %s not found in %s.", key, self.dotenv_path)
  94. return None
  95. def get_key(
  96. dotenv_path: StrPath,
  97. key_to_get: str,
  98. encoding: Optional[str] = "utf-8",
  99. ) -> Optional[str]:
  100. """
  101. Get the value of a given key from the given .env.
  102. Returns `None` if the key isn't found or doesn't have a value.
  103. """
  104. return DotEnv(dotenv_path, verbose=True, encoding=encoding).get(key_to_get)
  105. @contextmanager
  106. def rewrite(
  107. path: StrPath,
  108. encoding: Optional[str],
  109. ) -> Iterator[Tuple[IO[str], IO[str]]]:
  110. pathlib.Path(path).touch()
  111. with tempfile.NamedTemporaryFile(mode="w", encoding=encoding, delete=False) as dest:
  112. error = None
  113. try:
  114. with open(path, encoding=encoding) as source:
  115. yield (source, dest)
  116. except BaseException as err:
  117. error = err
  118. if error is None:
  119. shutil.move(dest.name, path)
  120. else:
  121. os.unlink(dest.name)
  122. raise error from None
  123. def set_key(
  124. dotenv_path: StrPath,
  125. key_to_set: str,
  126. value_to_set: str,
  127. quote_mode: str = "always",
  128. export: bool = False,
  129. encoding: Optional[str] = "utf-8",
  130. ) -> Tuple[Optional[bool], str, str]:
  131. """
  132. Adds or Updates a key/value to the given .env
  133. If the .env path given doesn't exist, fails instead of risking creating
  134. an orphan .env somewhere in the filesystem
  135. """
  136. if quote_mode not in ("always", "auto", "never"):
  137. raise ValueError(f"Unknown quote_mode: {quote_mode}")
  138. quote = quote_mode == "always" or (
  139. quote_mode == "auto" and not value_to_set.isalnum()
  140. )
  141. if quote:
  142. value_out = "'{}'".format(value_to_set.replace("'", "\\'"))
  143. else:
  144. value_out = value_to_set
  145. if export:
  146. line_out = f"export {key_to_set}={value_out}\n"
  147. else:
  148. line_out = f"{key_to_set}={value_out}\n"
  149. with rewrite(dotenv_path, encoding=encoding) as (source, dest):
  150. replaced = False
  151. missing_newline = False
  152. for mapping in with_warn_for_invalid_lines(parse_stream(source)):
  153. if mapping.key == key_to_set:
  154. dest.write(line_out)
  155. replaced = True
  156. else:
  157. dest.write(mapping.original.string)
  158. missing_newline = not mapping.original.string.endswith("\n")
  159. if not replaced:
  160. if missing_newline:
  161. dest.write("\n")
  162. dest.write(line_out)
  163. return True, key_to_set, value_to_set
  164. def unset_key(
  165. dotenv_path: StrPath,
  166. key_to_unset: str,
  167. quote_mode: str = "always",
  168. encoding: Optional[str] = "utf-8",
  169. ) -> Tuple[Optional[bool], str]:
  170. """
  171. Removes a given key from the given `.env` file.
  172. If the .env path given doesn't exist, fails.
  173. If the given key doesn't exist in the .env, fails.
  174. """
  175. if not os.path.exists(dotenv_path):
  176. logger.warning("Can't delete from %s - it doesn't exist.", dotenv_path)
  177. return None, key_to_unset
  178. removed = False
  179. with rewrite(dotenv_path, encoding=encoding) as (source, dest):
  180. for mapping in with_warn_for_invalid_lines(parse_stream(source)):
  181. if mapping.key == key_to_unset:
  182. removed = True
  183. else:
  184. dest.write(mapping.original.string)
  185. if not removed:
  186. logger.warning(
  187. "Key %s not removed from %s - key doesn't exist.", key_to_unset, dotenv_path
  188. )
  189. return None, key_to_unset
  190. return removed, key_to_unset
  191. def resolve_variables(
  192. values: Iterable[Tuple[str, Optional[str]]],
  193. override: bool,
  194. ) -> Mapping[str, Optional[str]]:
  195. new_values: Dict[str, Optional[str]] = {}
  196. for name, value in values:
  197. if value is None:
  198. result = None
  199. else:
  200. atoms = parse_variables(value)
  201. env: Dict[str, Optional[str]] = {}
  202. if override:
  203. env.update(os.environ) # type: ignore
  204. env.update(new_values)
  205. else:
  206. env.update(new_values)
  207. env.update(os.environ) # type: ignore
  208. result = "".join(atom.resolve(env) for atom in atoms)
  209. new_values[name] = result
  210. return new_values
  211. def _walk_to_root(path: str) -> Iterator[str]:
  212. """
  213. Yield directories starting from the given directory up to the root
  214. """
  215. if not os.path.exists(path):
  216. raise IOError("Starting path not found")
  217. if os.path.isfile(path):
  218. path = os.path.dirname(path)
  219. last_dir = None
  220. current_dir = os.path.abspath(path)
  221. while last_dir != current_dir:
  222. yield current_dir
  223. parent_dir = os.path.abspath(os.path.join(current_dir, os.path.pardir))
  224. last_dir, current_dir = current_dir, parent_dir
  225. def find_dotenv(
  226. filename: str = ".env",
  227. raise_error_if_not_found: bool = False,
  228. usecwd: bool = False,
  229. ) -> str:
  230. """
  231. Search in increasingly higher folders for the given file
  232. Returns path to the file if found, or an empty string otherwise
  233. """
  234. def _is_interactive():
  235. """Decide whether this is running in a REPL or IPython notebook"""
  236. if hasattr(sys, "ps1") or hasattr(sys, "ps2"):
  237. return True
  238. try:
  239. main = __import__("__main__", None, None, fromlist=["__file__"])
  240. except ModuleNotFoundError:
  241. return False
  242. return not hasattr(main, "__file__")
  243. def _is_debugger():
  244. return sys.gettrace() is not None
  245. if usecwd or _is_interactive() or _is_debugger() or getattr(sys, "frozen", False):
  246. # Should work without __file__, e.g. in REPL or IPython notebook.
  247. path = os.getcwd()
  248. else:
  249. # will work for .py files
  250. frame = sys._getframe()
  251. current_file = __file__
  252. while frame.f_code.co_filename == current_file or not os.path.exists(
  253. frame.f_code.co_filename
  254. ):
  255. assert frame.f_back is not None
  256. frame = frame.f_back
  257. frame_filename = frame.f_code.co_filename
  258. path = os.path.dirname(os.path.abspath(frame_filename))
  259. for dirname in _walk_to_root(path):
  260. check_path = os.path.join(dirname, filename)
  261. if os.path.isfile(check_path):
  262. return check_path
  263. if raise_error_if_not_found:
  264. raise IOError("File not found")
  265. return ""
  266. def load_dotenv(
  267. dotenv_path: Optional[StrPath] = None,
  268. stream: Optional[IO[str]] = None,
  269. verbose: bool = False,
  270. override: bool = False,
  271. interpolate: bool = True,
  272. encoding: Optional[str] = "utf-8",
  273. ) -> bool:
  274. """Parse a .env file and then load all the variables found as environment variables.
  275. Parameters:
  276. dotenv_path: Absolute or relative path to .env file.
  277. stream: Text stream (such as `io.StringIO`) with .env content, used if
  278. `dotenv_path` is `None`.
  279. verbose: Whether to output a warning the .env file is missing.
  280. override: Whether to override the system environment variables with the variables
  281. from the `.env` file.
  282. encoding: Encoding to be used to read the file.
  283. Returns:
  284. Bool: True if at least one environment variable is set else False
  285. If both `dotenv_path` and `stream` are `None`, `find_dotenv()` is used to find the
  286. .env file with it's default parameters. If you need to change the default parameters
  287. of `find_dotenv()`, you can explicitly call `find_dotenv()` and pass the result
  288. to this function as `dotenv_path`.
  289. """
  290. if dotenv_path is None and stream is None:
  291. dotenv_path = find_dotenv()
  292. dotenv = DotEnv(
  293. dotenv_path=dotenv_path,
  294. stream=stream,
  295. verbose=verbose,
  296. interpolate=interpolate,
  297. override=override,
  298. encoding=encoding,
  299. )
  300. return dotenv.set_as_environment_variables()
  301. def dotenv_values(
  302. dotenv_path: Optional[StrPath] = None,
  303. stream: Optional[IO[str]] = None,
  304. verbose: bool = False,
  305. interpolate: bool = True,
  306. encoding: Optional[str] = "utf-8",
  307. ) -> Dict[str, Optional[str]]:
  308. """
  309. Parse a .env file and return its content as a dict.
  310. The returned dict will have `None` values for keys without values in the .env file.
  311. For example, `foo=bar` results in `{"foo": "bar"}` whereas `foo` alone results in
  312. `{"foo": None}`
  313. Parameters:
  314. dotenv_path: Absolute or relative path to the .env file.
  315. stream: `StringIO` object with .env content, used if `dotenv_path` is `None`.
  316. verbose: Whether to output a warning if the .env file is missing.
  317. encoding: Encoding to be used to read the file.
  318. If both `dotenv_path` and `stream` are `None`, `find_dotenv()` is used to find the
  319. .env file.
  320. """
  321. if dotenv_path is None and stream is None:
  322. dotenv_path = find_dotenv()
  323. return DotEnv(
  324. dotenv_path=dotenv_path,
  325. stream=stream,
  326. verbose=verbose,
  327. interpolate=interpolate,
  328. override=True,
  329. encoding=encoding,
  330. ).dict()