codekingpro/portable-devtools

1import io2import logging3import os4import pathlib5import stat6import sys7import tempfile8from collections import OrderedDict9from contextlib import contextmanager10from typing import IO, Dict, Iterable, Iterator, Mapping, Optional, Tuple, Union11 12from .parser import Binding, parse_stream13from .variables import parse_variables14 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 `os.replace()`.17StrPath = Union[str, "os.PathLike[str]"]18 19logger = logging.getLogger(__name__)20 21 22def _load_dotenv_disabled() -> bool:23    """24    Determine if dotenv loading has been disabled.25    """26    if "PYTHON_DOTENV_DISABLED" not in os.environ:27        return False28    value = os.environ["PYTHON_DOTENV_DISABLED"].casefold()29    return value in {"1", "true", "t", "yes", "y"}30 31 32def with_warn_for_invalid_lines(mappings: Iterator[Binding]) -> Iterator[Binding]:33    for mapping in mappings:34        if mapping.error:35            logger.warning(36                "python-dotenv could not parse statement starting at line %s",37                mapping.original.line,38            )39        yield mapping40 41 42class DotEnv:43    def __init__(44        self,45        dotenv_path: Optional[StrPath],46        stream: Optional[IO[str]] = None,47        verbose: bool = False,48        encoding: Optional[str] = None,49        interpolate: bool = True,50        override: bool = True,51    ) -> None:52        self.dotenv_path: Optional[StrPath] = dotenv_path53        self.stream: Optional[IO[str]] = stream54        self._dict: Optional[Dict[str, Optional[str]]] = None55        self.verbose: bool = verbose56        self.encoding: Optional[str] = encoding57        self.interpolate: bool = interpolate58        self.override: bool = override59 60    @contextmanager61    def _get_stream(self) -> Iterator[IO[str]]:62        if self.dotenv_path and _is_file_or_fifo(self.dotenv_path):63            with open(self.dotenv_path, encoding=self.encoding) as stream:64                yield stream65        elif self.stream is not None:66            yield self.stream67        else:68            if self.verbose:69                logger.info(70                    "python-dotenv could not find configuration file %s.",71                    self.dotenv_path or ".env",72                )73            yield io.StringIO("")74 75    def dict(self) -> Dict[str, Optional[str]]:76        """Return dotenv as dict"""77        if self._dict:78            return self._dict79 80        raw_values = self.parse()81 82        if self.interpolate:83            self._dict = OrderedDict(84                resolve_variables(raw_values, override=self.override)85            )86        else:87            self._dict = OrderedDict(raw_values)88 89        return self._dict90 91    def parse(self) -> Iterator[Tuple[str, Optional[str]]]:92        with self._get_stream() as stream:93            for mapping in with_warn_for_invalid_lines(parse_stream(stream)):94                if mapping.key is not None:95                    yield mapping.key, mapping.value96 97    def set_as_environment_variables(self) -> bool:98        """99        Load the current dotenv as system environment variable.100        """101        if not self.dict():102            return False103 104        for k, v in self.dict().items():105            if k in os.environ and not self.override:106                continue107            if v is not None:108                os.environ[k] = v109 110        return True111 112    def get(self, key: str) -> Optional[str]:113        """ """114        data = self.dict()115 116        if key in data:117            return data[key]118 119        if self.verbose:120            logger.warning("Key %s not found in %s.", key, self.dotenv_path)121 122        return None123 124 125def get_key(126    dotenv_path: StrPath,127    key_to_get: str,128    encoding: Optional[str] = "utf-8",129) -> Optional[str]:130    """131    Get the value of a given key from the given .env.132 133    Returns `None` if the key isn't found or doesn't have a value.134    """135    return DotEnv(dotenv_path, verbose=True, encoding=encoding).get(key_to_get)136 137 138@contextmanager139def rewrite(140    path: StrPath,141    encoding: Optional[str],142    follow_symlinks: bool = False,143) -> Iterator[Tuple[IO[str], IO[str]]]:144    if follow_symlinks:145        path = os.path.realpath(path)146 147    try:148        source: IO[str] = open(path, encoding=encoding)149        try:150            path_stat = os.lstat(path)151            original_mode: Optional[int] = (152                stat.S_IMODE(path_stat.st_mode)153                if stat.S_ISREG(path_stat.st_mode)154                else None155            )156        except BaseException:157            source.close()158            raise159    except FileNotFoundError:160        source = io.StringIO("")161        original_mode = None162 163    with tempfile.NamedTemporaryFile(164        mode="w",165        encoding=encoding,166        delete=False,167        prefix=".tmp_",168        dir=os.path.dirname(os.path.abspath(path)),169    ) as dest:170        dest_path = pathlib.Path(dest.name)171        error = None172 173        try:174            with source:175                yield (source, dest)176        except BaseException as err:177            error = err178 179    if error is None:180        try:181            if original_mode is not None:182                os.chmod(dest_path, original_mode)183 184            os.replace(dest_path, path)185        except BaseException:186            dest_path.unlink(missing_ok=True)187            raise188    else:189        dest_path.unlink(missing_ok=True)190        raise error from None191 192 193def set_key(194    dotenv_path: StrPath,195    key_to_set: str,196    value_to_set: str,197    quote_mode: str = "always",198    export: bool = False,199    encoding: Optional[str] = "utf-8",200    follow_symlinks: bool = False,201) -> Tuple[Optional[bool], str, str]:202    """203    Adds or Updates a key/value to the given .env204 205    The target .env file is created if it doesn't exist.206 207    This function doesn't follow symlinks by default, to avoid accidentally208    modifying a file at a potentially untrusted path. If you don't need this209    protection and need symlinks to be followed, use `follow_symlinks`.210    """211    if quote_mode not in ("always", "auto", "never"):212        raise ValueError(f"Unknown quote_mode: {quote_mode}")213 214    quote = quote_mode == "always" or (215        quote_mode == "auto" and not value_to_set.isalnum()216    )217 218    if quote:219        value_out = "'{}'".format(value_to_set.replace("'", "\\'"))220    else:221        value_out = value_to_set222    if export:223        line_out = f"export {key_to_set}={value_out}\n"224    else:225        line_out = f"{key_to_set}={value_out}\n"226 227    with rewrite(dotenv_path, encoding=encoding, follow_symlinks=follow_symlinks) as (228        source,229        dest,230    ):231        replaced = False232        missing_newline = False233        for mapping in with_warn_for_invalid_lines(parse_stream(source)):234            if mapping.key == key_to_set:235                dest.write(line_out)236                replaced = True237            else:238                dest.write(mapping.original.string)239                missing_newline = not mapping.original.string.endswith("\n")240        if not replaced:241            if missing_newline:242                dest.write("\n")243            dest.write(line_out)244 245    return True, key_to_set, value_to_set246 247 248def unset_key(249    dotenv_path: StrPath,250    key_to_unset: str,251    quote_mode: str = "always",252    encoding: Optional[str] = "utf-8",253    follow_symlinks: bool = False,254) -> Tuple[Optional[bool], str]:255    """256    Removes a given key from the given `.env` file.257 258    If the .env path given doesn't exist, fails.259    If the given key doesn't exist in the .env, fails.260 261    This function doesn't follow symlinks by default, to avoid accidentally262    modifying a file at a potentially untrusted path. If you don't need this263    protection and need symlinks to be followed, use `follow_symlinks`.264    """265    if not os.path.exists(dotenv_path):266        logger.warning("Can't delete from %s - it doesn't exist.", dotenv_path)267        return None, key_to_unset268 269    removed = False270    with rewrite(dotenv_path, encoding=encoding, follow_symlinks=follow_symlinks) as (271        source,272        dest,273    ):274        for mapping in with_warn_for_invalid_lines(parse_stream(source)):275            if mapping.key == key_to_unset:276                removed = True277            else:278                dest.write(mapping.original.string)279 280    if not removed:281        logger.warning(282            "Key %s not removed from %s - key doesn't exist.", key_to_unset, dotenv_path283        )284        return None, key_to_unset285 286    return removed, key_to_unset287 288 289def resolve_variables(290    values: Iterable[Tuple[str, Optional[str]]],291    override: bool,292) -> Mapping[str, Optional[str]]:293    new_values: Dict[str, Optional[str]] = {}294 295    for name, value in values:296        if value is None:297            result = None298        else:299            atoms = parse_variables(value)300            env: Dict[str, Optional[str]] = {}301            if override:302                env.update(os.environ)  # type: ignore303                env.update(new_values)304            else:305                env.update(new_values)306                env.update(os.environ)  # type: ignore307            result = "".join(atom.resolve(env) for atom in atoms)308 309        new_values[name] = result310 311    return new_values312 313 314def _walk_to_root(path: str) -> Iterator[str]:315    """316    Yield directories starting from the given directory up to the root317    """318    if not os.path.exists(path):319        raise IOError("Starting path not found")320 321    if os.path.isfile(path):322        path = os.path.dirname(path)323 324    last_dir = None325    current_dir = os.path.abspath(path)326    while last_dir != current_dir:327        yield current_dir328        parent_dir = os.path.abspath(os.path.join(current_dir, os.path.pardir))329        last_dir, current_dir = current_dir, parent_dir330 331 332def find_dotenv(333    filename: str = ".env",334    raise_error_if_not_found: bool = False,335    usecwd: bool = False,336) -> str:337    """338    Search in increasingly higher folders for the given file339 340    Returns path to the file if found, or an empty string otherwise341    """342 343    def _is_interactive():344        """Decide whether this is running in a REPL or IPython notebook"""345        if hasattr(sys, "ps1") or hasattr(sys, "ps2"):346            return True347        try:348            main = __import__("__main__", None, None, fromlist=["__file__"])349        except ModuleNotFoundError:350            return False351        return not hasattr(main, "__file__")352 353    def _is_debugger():354        return sys.gettrace() is not None355 356    if usecwd or _is_interactive() or _is_debugger() or getattr(sys, "frozen", False):357        # Should work without __file__, e.g. in REPL or IPython notebook.358        path = os.getcwd()359    else:360        # will work for .py files361        frame = sys._getframe()362        current_file = __file__363 364        while frame.f_code.co_filename == current_file or not os.path.exists(365            frame.f_code.co_filename366        ):367            assert frame.f_back is not None368            frame = frame.f_back369        frame_filename = frame.f_code.co_filename370        path = os.path.dirname(os.path.abspath(frame_filename))371 372    for dirname in _walk_to_root(path):373        check_path = os.path.join(dirname, filename)374        if _is_file_or_fifo(check_path):375            return check_path376 377    if raise_error_if_not_found:378        raise IOError("File not found")379 380    return ""381 382 383def load_dotenv(384    dotenv_path: Optional[StrPath] = None,385    stream: Optional[IO[str]] = None,386    verbose: bool = False,387    override: bool = False,388    interpolate: bool = True,389    encoding: Optional[str] = "utf-8",390) -> bool:391    """Parse a .env file and then load all the variables found as environment variables.392 393    Parameters:394        dotenv_path: Absolute or relative path to .env file.395        stream: Text stream (such as `io.StringIO`) with .env content, used if396            `dotenv_path` is `None`.397        verbose: Whether to output a warning the .env file is missing.398        override: Whether to override the system environment variables with the variables399            from the `.env` file.400        encoding: Encoding to be used to read the file.401    Returns:402        Bool: True if at least one environment variable is set else False403 404    If both `dotenv_path` and `stream` are `None`, `find_dotenv()` is used to find the405    .env file with it's default parameters. If you need to change the default parameters406    of `find_dotenv()`, you can explicitly call `find_dotenv()` and pass the result407    to this function as `dotenv_path`.408 409    If the environment variable `PYTHON_DOTENV_DISABLED` is set to a truthy value,410    .env loading is disabled.411    """412    if _load_dotenv_disabled():413        logger.debug(414            "python-dotenv: .env loading disabled by PYTHON_DOTENV_DISABLED environment variable"415        )416        return False417 418    if dotenv_path is None and stream is None:419        dotenv_path = find_dotenv()420 421    dotenv = DotEnv(422        dotenv_path=dotenv_path,423        stream=stream,424        verbose=verbose,425        interpolate=interpolate,426        override=override,427        encoding=encoding,428    )429    return dotenv.set_as_environment_variables()430 431 432def dotenv_values(433    dotenv_path: Optional[StrPath] = None,434    stream: Optional[IO[str]] = None,435    verbose: bool = False,436    interpolate: bool = True,437    encoding: Optional[str] = "utf-8",438) -> Dict[str, Optional[str]]:439    """440    Parse a .env file and return its content as a dict.441 442    The returned dict will have `None` values for keys without values in the .env file.443    For example, `foo=bar` results in `{"foo": "bar"}` whereas `foo` alone results in444    `{"foo": None}`445 446    Parameters:447        dotenv_path: Absolute or relative path to the .env file.448        stream: `StringIO` object with .env content, used if `dotenv_path` is `None`.449        verbose: Whether to output a warning if the .env file is missing.450        encoding: Encoding to be used to read the file.451 452    If both `dotenv_path` and `stream` are `None`, `find_dotenv()` is used to find the453    .env file.454    """455    if dotenv_path is None and stream is None:456        dotenv_path = find_dotenv()457 458    return DotEnv(459        dotenv_path=dotenv_path,460        stream=stream,461        verbose=verbose,462        interpolate=interpolate,463        override=True,464        encoding=encoding,465    ).dict()466 467 468def _is_file_or_fifo(path: StrPath) -> bool:469    """470    Return True if `path` exists and is either a regular file or a FIFO.471    """472    if os.path.isfile(path):473        return True474 475    try:476        st = os.stat(path)477    except (FileNotFoundError, OSError):478        return False479 480    return stat.S_ISFIFO(st.st_mode)481