main.py 14 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480
  1. import io
  2. import logging
  3. import os
  4. import pathlib
  5. import stat
  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 `os.replace()`.
  15. StrPath = Union[str, "os.PathLike[str]"]
  16. logger = logging.getLogger(__name__)
  17. def _load_dotenv_disabled() -> bool:
  18. """
  19. Determine if dotenv loading has been disabled.
  20. """
  21. if "PYTHON_DOTENV_DISABLED" not in os.environ:
  22. return False
  23. value = os.environ["PYTHON_DOTENV_DISABLED"].casefold()
  24. return value in {"1", "true", "t", "yes", "y"}
  25. def with_warn_for_invalid_lines(mappings: Iterator[Binding]) -> Iterator[Binding]:
  26. for mapping in mappings:
  27. if mapping.error:
  28. logger.warning(
  29. "python-dotenv could not parse statement starting at line %s",
  30. mapping.original.line,
  31. )
  32. yield mapping
  33. class DotEnv:
  34. def __init__(
  35. self,
  36. dotenv_path: Optional[StrPath],
  37. stream: Optional[IO[str]] = None,
  38. verbose: bool = False,
  39. encoding: Optional[str] = None,
  40. interpolate: bool = True,
  41. override: bool = True,
  42. ) -> None:
  43. self.dotenv_path: Optional[StrPath] = dotenv_path
  44. self.stream: Optional[IO[str]] = stream
  45. self._dict: Optional[Dict[str, Optional[str]]] = None
  46. self.verbose: bool = verbose
  47. self.encoding: Optional[str] = encoding
  48. self.interpolate: bool = interpolate
  49. self.override: bool = override
  50. @contextmanager
  51. def _get_stream(self) -> Iterator[IO[str]]:
  52. if self.dotenv_path and _is_file_or_fifo(self.dotenv_path):
  53. with open(self.dotenv_path, encoding=self.encoding) as stream:
  54. yield stream
  55. elif self.stream is not None:
  56. yield self.stream
  57. else:
  58. if self.verbose:
  59. logger.info(
  60. "python-dotenv could not find configuration file %s.",
  61. self.dotenv_path or ".env",
  62. )
  63. yield io.StringIO("")
  64. def dict(self) -> Dict[str, Optional[str]]:
  65. """Return dotenv as dict"""
  66. if self._dict:
  67. return self._dict
  68. raw_values = self.parse()
  69. if self.interpolate:
  70. self._dict = OrderedDict(
  71. resolve_variables(raw_values, override=self.override)
  72. )
  73. else:
  74. self._dict = OrderedDict(raw_values)
  75. return self._dict
  76. def parse(self) -> Iterator[Tuple[str, Optional[str]]]:
  77. with self._get_stream() as stream:
  78. for mapping in with_warn_for_invalid_lines(parse_stream(stream)):
  79. if mapping.key is not None:
  80. yield mapping.key, mapping.value
  81. def set_as_environment_variables(self) -> bool:
  82. """
  83. Load the current dotenv as system environment variable.
  84. """
  85. if not self.dict():
  86. return False
  87. for k, v in self.dict().items():
  88. if k in os.environ and not self.override:
  89. continue
  90. if v is not None:
  91. os.environ[k] = v
  92. return True
  93. def get(self, key: str) -> Optional[str]:
  94. """ """
  95. data = self.dict()
  96. if key in data:
  97. return data[key]
  98. if self.verbose:
  99. logger.warning("Key %s not found in %s.", key, self.dotenv_path)
  100. return None
  101. def get_key(
  102. dotenv_path: StrPath,
  103. key_to_get: str,
  104. encoding: Optional[str] = "utf-8",
  105. ) -> Optional[str]:
  106. """
  107. Get the value of a given key from the given .env.
  108. Returns `None` if the key isn't found or doesn't have a value.
  109. """
  110. return DotEnv(dotenv_path, verbose=True, encoding=encoding).get(key_to_get)
  111. @contextmanager
  112. def rewrite(
  113. path: StrPath,
  114. encoding: Optional[str],
  115. follow_symlinks: bool = False,
  116. ) -> Iterator[Tuple[IO[str], IO[str]]]:
  117. if follow_symlinks:
  118. path = os.path.realpath(path)
  119. try:
  120. source: IO[str] = open(path, encoding=encoding)
  121. try:
  122. path_stat = os.lstat(path)
  123. original_mode: Optional[int] = (
  124. stat.S_IMODE(path_stat.st_mode)
  125. if stat.S_ISREG(path_stat.st_mode)
  126. else None
  127. )
  128. except BaseException:
  129. source.close()
  130. raise
  131. except FileNotFoundError:
  132. source = io.StringIO("")
  133. original_mode = None
  134. with tempfile.NamedTemporaryFile(
  135. mode="w",
  136. encoding=encoding,
  137. delete=False,
  138. prefix=".tmp_",
  139. dir=os.path.dirname(os.path.abspath(path)),
  140. ) as dest:
  141. dest_path = pathlib.Path(dest.name)
  142. error = None
  143. try:
  144. with source:
  145. yield (source, dest)
  146. except BaseException as err:
  147. error = err
  148. if error is None:
  149. try:
  150. if original_mode is not None:
  151. os.chmod(dest_path, original_mode)
  152. os.replace(dest_path, path)
  153. except BaseException:
  154. dest_path.unlink(missing_ok=True)
  155. raise
  156. else:
  157. dest_path.unlink(missing_ok=True)
  158. raise error from None
  159. def set_key(
  160. dotenv_path: StrPath,
  161. key_to_set: str,
  162. value_to_set: str,
  163. quote_mode: str = "always",
  164. export: bool = False,
  165. encoding: Optional[str] = "utf-8",
  166. follow_symlinks: bool = False,
  167. ) -> Tuple[Optional[bool], str, str]:
  168. """
  169. Adds or Updates a key/value to the given .env
  170. The target .env file is created if it doesn't exist.
  171. This function doesn't follow symlinks by default, to avoid accidentally
  172. modifying a file at a potentially untrusted path. If you don't need this
  173. protection and need symlinks to be followed, use `follow_symlinks`.
  174. """
  175. if quote_mode not in ("always", "auto", "never"):
  176. raise ValueError(f"Unknown quote_mode: {quote_mode}")
  177. quote = quote_mode == "always" or (
  178. quote_mode == "auto" and not value_to_set.isalnum()
  179. )
  180. if quote:
  181. value_out = "'{}'".format(value_to_set.replace("'", "\\'"))
  182. else:
  183. value_out = value_to_set
  184. if export:
  185. line_out = f"export {key_to_set}={value_out}\n"
  186. else:
  187. line_out = f"{key_to_set}={value_out}\n"
  188. with rewrite(dotenv_path, encoding=encoding, follow_symlinks=follow_symlinks) as (
  189. source,
  190. dest,
  191. ):
  192. replaced = False
  193. missing_newline = False
  194. for mapping in with_warn_for_invalid_lines(parse_stream(source)):
  195. if mapping.key == key_to_set:
  196. dest.write(line_out)
  197. replaced = True
  198. else:
  199. dest.write(mapping.original.string)
  200. missing_newline = not mapping.original.string.endswith("\n")
  201. if not replaced:
  202. if missing_newline:
  203. dest.write("\n")
  204. dest.write(line_out)
  205. return True, key_to_set, value_to_set
  206. def unset_key(
  207. dotenv_path: StrPath,
  208. key_to_unset: str,
  209. quote_mode: str = "always",
  210. encoding: Optional[str] = "utf-8",
  211. follow_symlinks: bool = False,
  212. ) -> Tuple[Optional[bool], str]:
  213. """
  214. Removes a given key from the given `.env` file.
  215. If the .env path given doesn't exist, fails.
  216. If the given key doesn't exist in the .env, fails.
  217. This function doesn't follow symlinks by default, to avoid accidentally
  218. modifying a file at a potentially untrusted path. If you don't need this
  219. protection and need symlinks to be followed, use `follow_symlinks`.
  220. """
  221. if not os.path.exists(dotenv_path):
  222. logger.warning("Can't delete from %s - it doesn't exist.", dotenv_path)
  223. return None, key_to_unset
  224. removed = False
  225. with rewrite(dotenv_path, encoding=encoding, follow_symlinks=follow_symlinks) as (
  226. source,
  227. dest,
  228. ):
  229. for mapping in with_warn_for_invalid_lines(parse_stream(source)):
  230. if mapping.key == key_to_unset:
  231. removed = True
  232. else:
  233. dest.write(mapping.original.string)
  234. if not removed:
  235. logger.warning(
  236. "Key %s not removed from %s - key doesn't exist.", key_to_unset, dotenv_path
  237. )
  238. return None, key_to_unset
  239. return removed, key_to_unset
  240. def resolve_variables(
  241. values: Iterable[Tuple[str, Optional[str]]],
  242. override: bool,
  243. ) -> Mapping[str, Optional[str]]:
  244. new_values: Dict[str, Optional[str]] = {}
  245. for name, value in values:
  246. if value is None:
  247. result = None
  248. else:
  249. atoms = parse_variables(value)
  250. env: Dict[str, Optional[str]] = {}
  251. if override:
  252. env.update(os.environ) # type: ignore
  253. env.update(new_values)
  254. else:
  255. env.update(new_values)
  256. env.update(os.environ) # type: ignore
  257. result = "".join(atom.resolve(env) for atom in atoms)
  258. new_values[name] = result
  259. return new_values
  260. def _walk_to_root(path: str) -> Iterator[str]:
  261. """
  262. Yield directories starting from the given directory up to the root
  263. """
  264. if not os.path.exists(path):
  265. raise IOError("Starting path not found")
  266. if os.path.isfile(path):
  267. path = os.path.dirname(path)
  268. last_dir = None
  269. current_dir = os.path.abspath(path)
  270. while last_dir != current_dir:
  271. yield current_dir
  272. parent_dir = os.path.abspath(os.path.join(current_dir, os.path.pardir))
  273. last_dir, current_dir = current_dir, parent_dir
  274. def find_dotenv(
  275. filename: str = ".env",
  276. raise_error_if_not_found: bool = False,
  277. usecwd: bool = False,
  278. ) -> str:
  279. """
  280. Search in increasingly higher folders for the given file
  281. Returns path to the file if found, or an empty string otherwise
  282. """
  283. def _is_interactive():
  284. """Decide whether this is running in a REPL or IPython notebook"""
  285. if hasattr(sys, "ps1") or hasattr(sys, "ps2"):
  286. return True
  287. try:
  288. main = __import__("__main__", None, None, fromlist=["__file__"])
  289. except ModuleNotFoundError:
  290. return False
  291. return not hasattr(main, "__file__")
  292. def _is_debugger():
  293. return sys.gettrace() is not None
  294. if usecwd or _is_interactive() or _is_debugger() or getattr(sys, "frozen", False):
  295. # Should work without __file__, e.g. in REPL or IPython notebook.
  296. path = os.getcwd()
  297. else:
  298. # will work for .py files
  299. frame = sys._getframe()
  300. current_file = __file__
  301. while frame.f_code.co_filename == current_file or not os.path.exists(
  302. frame.f_code.co_filename
  303. ):
  304. assert frame.f_back is not None
  305. frame = frame.f_back
  306. frame_filename = frame.f_code.co_filename
  307. path = os.path.dirname(os.path.abspath(frame_filename))
  308. for dirname in _walk_to_root(path):
  309. check_path = os.path.join(dirname, filename)
  310. if _is_file_or_fifo(check_path):
  311. return check_path
  312. if raise_error_if_not_found:
  313. raise IOError("File not found")
  314. return ""
  315. def load_dotenv(
  316. dotenv_path: Optional[StrPath] = None,
  317. stream: Optional[IO[str]] = None,
  318. verbose: bool = False,
  319. override: bool = False,
  320. interpolate: bool = True,
  321. encoding: Optional[str] = "utf-8",
  322. ) -> bool:
  323. """Parse a .env file and then load all the variables found as environment variables.
  324. Parameters:
  325. dotenv_path: Absolute or relative path to .env file.
  326. stream: Text stream (such as `io.StringIO`) with .env content, used if
  327. `dotenv_path` is `None`.
  328. verbose: Whether to output a warning the .env file is missing.
  329. override: Whether to override the system environment variables with the variables
  330. from the `.env` file.
  331. encoding: Encoding to be used to read the file.
  332. Returns:
  333. Bool: True if at least one environment variable is set else False
  334. If both `dotenv_path` and `stream` are `None`, `find_dotenv()` is used to find the
  335. .env file with it's default parameters. If you need to change the default parameters
  336. of `find_dotenv()`, you can explicitly call `find_dotenv()` and pass the result
  337. to this function as `dotenv_path`.
  338. If the environment variable `PYTHON_DOTENV_DISABLED` is set to a truthy value,
  339. .env loading is disabled.
  340. """
  341. if _load_dotenv_disabled():
  342. logger.debug(
  343. "python-dotenv: .env loading disabled by PYTHON_DOTENV_DISABLED environment variable"
  344. )
  345. return False
  346. if dotenv_path is None and stream is None:
  347. dotenv_path = find_dotenv()
  348. dotenv = DotEnv(
  349. dotenv_path=dotenv_path,
  350. stream=stream,
  351. verbose=verbose,
  352. interpolate=interpolate,
  353. override=override,
  354. encoding=encoding,
  355. )
  356. return dotenv.set_as_environment_variables()
  357. def dotenv_values(
  358. dotenv_path: Optional[StrPath] = None,
  359. stream: Optional[IO[str]] = None,
  360. verbose: bool = False,
  361. interpolate: bool = True,
  362. encoding: Optional[str] = "utf-8",
  363. ) -> Dict[str, Optional[str]]:
  364. """
  365. Parse a .env file and return its content as a dict.
  366. The returned dict will have `None` values for keys without values in the .env file.
  367. For example, `foo=bar` results in `{"foo": "bar"}` whereas `foo` alone results in
  368. `{"foo": None}`
  369. Parameters:
  370. dotenv_path: Absolute or relative path to the .env file.
  371. stream: `StringIO` object with .env content, used if `dotenv_path` is `None`.
  372. verbose: Whether to output a warning if the .env file is missing.
  373. encoding: Encoding to be used to read the file.
  374. If both `dotenv_path` and `stream` are `None`, `find_dotenv()` is used to find the
  375. .env file.
  376. """
  377. if dotenv_path is None and stream is None:
  378. dotenv_path = find_dotenv()
  379. return DotEnv(
  380. dotenv_path=dotenv_path,
  381. stream=stream,
  382. verbose=verbose,
  383. interpolate=interpolate,
  384. override=True,
  385. encoding=encoding,
  386. ).dict()
  387. def _is_file_or_fifo(path: StrPath) -> bool:
  388. """
  389. Return True if `path` exists and is either a regular file or a FIFO.
  390. """
  391. if os.path.isfile(path):
  392. return True
  393. try:
  394. st = os.stat(path)
  395. except (FileNotFoundError, OSError):
  396. return False
  397. return stat.S_ISFIFO(st.st_mode)