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