From a82faf91f414c2387fa5d75687d6ce7e3dfbb4fa Mon Sep 17 00:00:00 2001
From: Ulysses Souza <ulyssessouza@gmail.com>
Date: Fri, 13 Mar 2020 18:08:34 +0100
Subject: [PATCH] Add POSIX default value support on vars

Signed-off-by: Ulysses Souza <ulyssessouza@gmail.com>
---
 src/dotenv/main.py   | 39 ++++++++++++++-------------------------
 src/dotenv/parser.py | 17 ++++++++++++++---
 tests/test_main.py   |  8 +++++++-
 3 files changed, 35 insertions(+), 29 deletions(-)

diff --git a/src/dotenv/main.py b/src/dotenv/main.py
index 93d617d6..44bd169a 100644
--- a/src/dotenv/main.py
+++ b/src/dotenv/main.py
@@ -12,13 +12,13 @@
 from contextlib import contextmanager
 
 from .compat import IS_TYPE_CHECKING, PY2, StringIO, to_env
-from .parser import Binding, parse_stream
+from .parser import Binding, Reader, parse_stream, parse_var_value, parse_value
 
 logger = logging.getLogger(__name__)
 
 if IS_TYPE_CHECKING:
     from typing import (
-        Dict, Iterator, Match, Optional, Pattern, Union, Text, IO, Tuple
+        Dict, Iterator, Optional, Pattern, Union, Text, IO, Tuple
     )
     if sys.version_info >= (3, 6):
         _PathLike = os.PathLike
@@ -53,6 +53,7 @@ def __init__(self, dotenv_path, verbose=False, encoding=None, interpolate=True):
         self.verbose = verbose  # type: bool
         self.encoding = encoding  # type: Union[None, Text]
         self.interpolate = interpolate  # type: bool
+        self.local_env = {}  # type: Dict[Text, Optional[Text]]
 
     @contextmanager
     def _get_stream(self):
@@ -200,32 +201,20 @@ def unset_key(dotenv_path, key_to_unset, quote_mode="always"):
     return removed, key_to_unset
 
 
-def resolve_nested_variables(values):
-    # type: (Dict[Text, Optional[Text]]) -> Dict[Text, Optional[Text]]
-    def _replacement(name):
-        # type: (Text) -> Text
-        """
-        get appropriate value for a variable name.
-        first search in environ, if not found,
-        then look into the dotenv variables
-        """
-        ret = os.getenv(name, new_values.get(name, ""))
-        return ret  # type: ignore
-
-    def _re_sub_callback(match_object):
-        # type: (Match[Text]) -> Text
-        """
-        From a match object gets the variable name and returns
-        the correct replacement
-        """
-        return _replacement(match_object.group()[2:-1])
-
-    new_values = {}
+def resolve_nested_variables(values, local_env=None):
+    # type: (Dict[Text, Optional[Text]], Dict[Text, Optional[Text]]) -> Dict[Text, Optional[Text]]
+    local_env = local_env or {}
 
     for k, v in values.items():
-        new_values[k] = __posix_variable.sub(_re_sub_callback, v) if v is not None else None
+        if not v:
+            value = v
+        elif v[0] == '$':
+            value = parse_var_value(Reader(StringIO(v)), local_env=local_env)
+        else:
+            value = parse_value(Reader(StringIO(v)))
+        local_env[k] = value
 
-    return new_values
+    return local_env
 
 
 def _walk_to_root(path):
diff --git a/src/dotenv/parser.py b/src/dotenv/parser.py
index 2c93cbd0..02e51ac5 100644
--- a/src/dotenv/parser.py
+++ b/src/dotenv/parser.py
@@ -1,13 +1,13 @@
 import codecs
+import os
 import re
 
 from .compat import IS_TYPE_CHECKING, to_text
 
 if IS_TYPE_CHECKING:
     from typing import (  # noqa:F401
-        IO, Iterator, Match, NamedTuple, Optional, Pattern, Sequence, Text,
-        Tuple
-    )
+    IO, Iterator, Match, NamedTuple, Optional, Pattern, Sequence, Text,
+    Tuple, Dict)
 
 
 def make_regex(string, extra_flags=0):
@@ -24,6 +24,7 @@ def make_regex(string, extra_flags=0):
 _equal_sign = make_regex(r"(=[^\S\r\n]*)")
 _single_quoted_value = make_regex(r"'((?:\\'|[^'])*)'")
 _double_quoted_value = make_regex(r'"((?:\\"|[^"])*)"')
+_variable_value = make_regex(r"\$({)?([a-zA-Z_]+)(:-[a-zA-Z_]+)?(})?")
 _unquoted_value_part = make_regex(r"([^ \r\n]*)")
 _comment = make_regex(r"(?:[^\S\r\n]*#[^\r\n]*)?")
 _end_of_line = make_regex(r"[^\S\r\n]*(?:\r\n|\n|\r|$)")
@@ -177,6 +178,16 @@ def parse_unquoted_value(reader):
         value += reader.read(2)
 
 
+def parse_var_value(reader, local_env=None):
+    # type: (Reader, Dict[Text, Optional[Text]]) -> Optional[Text]
+    env = local_env or {}
+    (_, var, default_val, _) = reader.read_regex(_variable_value)
+    default = default_val if default_val else ''
+    if default:
+        default = default[2:]
+    return env.get(var, os.environ.get(str(var), default))
+
+
 def parse_value(reader):
     # type: (Reader) -> Text
     char = reader.peek(1)
diff --git a/tests/test_main.py b/tests/test_main.py
index d8678589..ea93785c 100644
--- a/tests/test_main.py
+++ b/tests/test_main.py
@@ -301,7 +301,7 @@ def test_dotenv_values_file(dotenv_file):
     [
         # Defined in environment, with and without interpolation
         ({"b": "c"}, "a=$b", False, {"a": "$b"}),
-        ({"b": "c"}, "a=$b", True, {"a": "$b"}),
+        ({"b": "c"}, "a=$b", True, {"a": "c"}),
         ({"b": "c"}, "a=${b}", False, {"a": "${b}"}),
         ({"b": "c"}, "a=${b}", True, {"a": "c"}),
 
@@ -311,6 +311,12 @@ def test_dotenv_values_file(dotenv_file):
         # Undefined
         ({}, "a=${b}", True, {"a": ""}),
 
+        # Undefined with default value
+        ({}, "a=${b:-DEFAULT_VALUE}", False, {"a": "${b:-DEFAULT_VALUE}"}),
+        ({}, "a=${b:-DEFAULT_VALUE}", True, {"a": "DEFAULT_VALUE"}),
+        ({"b": "c"}, "a=${b:-DEFAULT_VALUE}", False, {"a": "${b:-DEFAULT_VALUE}"}),
+        ({"b": "c"}, "a=${b:-DEFAULT_VALUE}", True, {"a": "c"}),
+
         # With quotes
         ({"b": "c"}, 'a="${b}"', True, {"a": "c"}),
         ({"b": "c"}, "a='${b}'", True, {"a": "c"}),