From c34e1a102f395df3625f6f2c69c6f062f1b5aa1a Mon Sep 17 00:00:00 2001 From: S-P Chan Date: Sat, 16 Aug 2025 10:47:38 +0800 Subject: [PATCH] Add option to restrict POSIX variable name regex --- src/dotenv/__init__.py | 11 ++++++++++- src/dotenv/main.py | 13 ++++++++++++- src/dotenv/variables.py | 21 ++++++++++++++++++--- tests/test_variables.py | 23 ++++++++++++++++++++++- 4 files changed, 62 insertions(+), 6 deletions(-) diff --git a/src/dotenv/__init__.py b/src/dotenv/__init__.py index dde24a01..d05299ba 100644 --- a/src/dotenv/__init__.py +++ b/src/dotenv/__init__.py @@ -1,6 +1,14 @@ from typing import Any, Optional -from .main import dotenv_values, find_dotenv, get_key, load_dotenv, set_key, unset_key +from .main import ( + dotenv_values, + find_dotenv, + get_key, + load_dotenv, + set_key, + set_variable_name_pattern, + unset_key, +) def load_ipython_extension(ipython: Any) -> None: @@ -48,4 +56,5 @@ def get_cli_string( "unset_key", "find_dotenv", "load_ipython_extension", + "set_variable_name_pattern", ] diff --git a/src/dotenv/main.py b/src/dotenv/main.py index b6de171c..0d469629 100644 --- a/src/dotenv/main.py +++ b/src/dotenv/main.py @@ -10,7 +10,7 @@ from typing import IO, Dict, Iterable, Iterator, Mapping, Optional, Tuple, Union from .parser import Binding, parse_stream -from .variables import parse_variables +from .variables import parse_variables, set_variable_name_pattern # A type alias for a string path to be used for the paths in this file. # These paths may flow to `open()` and `shutil.move()`; `shutil.move()` @@ -341,6 +341,7 @@ def load_dotenv( override: bool = False, interpolate: bool = True, encoding: Optional[str] = "utf-8", + varname_pattern: Optional[str] = None, ) -> bool: """Parse a .env file and then load all the variables found as environment variables. @@ -352,6 +353,8 @@ def load_dotenv( override: Whether to override the system environment variables with the variables from the `.env` file. encoding: Encoding to be used to read the file. + varname_pattern: Optional regex pattern to restrict variable names. + If `None`, the existing pattern is used. The pattern set here is persistent. Returns: Bool: True if at least one environment variable is set else False @@ -380,6 +383,9 @@ def load_dotenv( override=override, encoding=encoding, ) + + if varname_pattern is not None: + set_variable_name_pattern(varname_pattern) return dotenv.set_as_environment_variables() @@ -389,6 +395,7 @@ def dotenv_values( verbose: bool = False, interpolate: bool = True, encoding: Optional[str] = "utf-8", + varname_pattern: Optional[str] = None, ) -> Dict[str, Optional[str]]: """ Parse a .env file and return its content as a dict. @@ -402,6 +409,8 @@ def dotenv_values( stream: `StringIO` object with .env content, used if `dotenv_path` is `None`. verbose: Whether to output a warning if the .env file is missing. encoding: Encoding to be used to read the file. + varname_pattern: Optional regex pattern to restrict variable names. + If `None`, the existing pattern is used. The pattern set here is persistent. If both `dotenv_path` and `stream` are `None`, `find_dotenv()` is used to find the .env file. @@ -409,6 +418,8 @@ def dotenv_values( if dotenv_path is None and stream is None: dotenv_path = find_dotenv() + if varname_pattern is not None: + set_variable_name_pattern(varname_pattern) return DotEnv( dotenv_path=dotenv_path, stream=stream, diff --git a/src/dotenv/variables.py b/src/dotenv/variables.py index 667f2f26..2762ed72 100644 --- a/src/dotenv/variables.py +++ b/src/dotenv/variables.py @@ -1,11 +1,14 @@ import re from abc import ABCMeta, abstractmethod -from typing import Iterator, Mapping, Optional, Pattern +from typing import Iterator, Mapping, Optional -_posix_variable: Pattern[str] = re.compile( +DEFAULT_VARNAME_RE = r"""[^\}:]*""" +def _pattern_builder(pattern: Optional[str] = None) -> re.Pattern[str]: + """Builds a regex pattern for ${xxx:-yyy} variable substitution.""" + return re.compile( r""" \$\{ - (?P[^\}:]*) + (?P""" + (pattern if pattern else DEFAULT_VARNAME_RE) + r""") (?::- (?P[^\}]*) )? @@ -15,6 +18,18 @@ ) +_posix_variable = _pattern_builder() + + +def set_variable_name_pattern(pattern: Optional[str] = None) -> None: + """Set the variable name pattern used by `parse_variables`. + + If `pattern` is None, it resets to the default pattern. + """ + global _posix_variable + _posix_variable = _pattern_builder(pattern) + + class Atom(metaclass=ABCMeta): def __ne__(self, other: object) -> bool: result = self.__eq__(other) diff --git a/tests/test_variables.py b/tests/test_variables.py index 6f2b2203..d4303504 100644 --- a/tests/test_variables.py +++ b/tests/test_variables.py @@ -1,6 +1,11 @@ import pytest -from dotenv.variables import Literal, Variable, parse_variables +from dotenv.variables import ( + Literal, + Variable, + parse_variables, + set_variable_name_pattern, +) @pytest.mark.parametrize( @@ -33,3 +38,19 @@ def test_parse_variables(value, expected): result = parse_variables(value) assert list(result) == expected + +@pytest.mark.parametrize( + "value,expected", + [ + ("", []), + ("${AB_CD}", [Variable(name="AB_CD", default=None)]), + ("${A.B.C.D}", [Literal(value="${A.B.C.D}")]), + ("${a}", [Literal(value="${a}")]), + ], +) +def test_parse_variables_re(value, expected): + set_variable_name_pattern(r"""[A-Z0-9_]+""") + result = parse_variables(value) + + assert list(result) == expected + set_variable_name_pattern(None)