Improve environment variable expansion for python subprocess

This makes it so `Config.launch` should resolve the same environment variables as the `hab launch` cli.

- Convert Formatters expand boolean into a mode Enum. It can now also replace
undefined `!e` variables with nothing like shells do.
- `Config.launch` uses the remove mode to replicate the shell.

Author: MHendricks

hab/formatter.py

@@ -1,9 +1,28 @@
+import enum
 import os
 import string
 
 from . import utils
 
 
+class ExpandMode(enum.Enum):
+    """Defines how Formatter replaces `!r` conformer flags."""
+
+    Preserve = (1, False)
+    """Preserve `!e` in Formatter calls. ('-{A!e}-' -> `-{A!e}-`)."""
+    ToShell = (2, True)
+    """Convert `!e` to the shell's environment variable. ('-{A!e}-' -> `-$A-`)."""
+    Remove = (3, True)
+    """Replace `!e` with a empty string. ('-{A!e}-' -> `--`)."""
+
+    def __init__(self, value, expand) -> None:
+        self._value_ = value
+        self.expand = expand
+
+    def __str__(self):
+        return self.name
+
+
 class Formatter(string.Formatter):
     """A extended string formatter class to parse habitat configurations.
 
@@ -23,8 +42,11 @@ class Formatter(string.Formatter):
             :py:const:`Formatter.shell_formats` and :py:meth:`Formatter.language_from_ext`.
             for supported values. If you pass None, it will preserve the
             formatting markers so it can be converted later.
-        expand(bool, optional): Should the ``!e`` conversion insert the shell
-            environment variable specifier or the value of the env var?
+        expand(ExpandMode, optional): Controls how the ``!e`` conversion is handled.
+        hab_env(dict, optional): If passed a dict of environment variable values
+            defined and updated by hab. These are used to replace `!e` vars. If
+            not found then falls back to os.environ values. This is used to expand
+            a previously defined hab variable into another hab variable.
 
     .. _conversion field:
        https://docs.python.org/3/library/string.html#grammar-token-format-string-conversion
@@ -64,10 +86,24 @@ class Formatter(string.Formatter):
     string that accepts the env var name. ``;`` is the path separator to use.
     """
 
-    def __init__(self, language, expand=False):
+    def __init__(self, language, expand=ExpandMode.Preserve, hab_env=None):
         super().__init__()
         self.language = self.language_from_ext(language)
         self.expand = expand
+        self.hab_env = {} if hab_env is None else hab_env
+
+    def format(self, format_string, *args, hab_env=None, **kwargs):
+        if hab_env is not None:
+            # Use the hab_env defined values instead of os.environ if possible.
+            try:
+                current = self.hab_env
+                self.hab_env = hab_env
+                return super().format(format_string, *args, **kwargs)
+            finally:
+                self.hab_env = current
+        else:
+            # fallback to the os.environ and other defined kwargs.
+            return super().format(format_string, *args, **kwargs)
 
     def get_field(self, field_name, args, kwargs):
         """Returns the object to be inserted for the given field_name.
@@ -80,6 +116,13 @@ def get_field(self, field_name, args, kwargs):
         .. _`string.Formatter`:
            https://docs.python.org/3/library/string.html#string.Formatter.get_field
         """
+        if field_name not in kwargs and field_name in self.hab_env:
+            # if field_name was not provided, use value stored in self.hab_env
+            value = self.hab_env[field_name]
+            if isinstance(value, list) and len(value):
+                value = utils.Platform.collapse_paths(value, ext=self.language)
+            return value, field_name
+
         # If a field_name was not provided, use the value stored in os.environ
         if field_name not in kwargs and field_name in os.environ:
             return os.getenv(field_name), field_name
@@ -120,12 +163,20 @@ def parse(self, txt):
                 yield (literal_text, field_name, format_spec, conversion)
                 continue
 
-            elif self.expand and field_name in os.environ:
+            elif self.expand.expand:
                 # Expand the env var to the env var value. Later `get_field`
-                # will update kwargs with the existing env var value
-                yield (literal_text, field_name, format_spec, "s")
-                continue
+                if field_name in self.hab_env:
+                    yield (literal_text, field_name, format_spec, "s")
+                    continue
+                elif field_name in os.environ:
+                    # Expand the env var to the env var value. Later `get_field`
+                    # will update kwargs with the existing env var value
+                    yield (literal_text, field_name, format_spec, "s")
+                    continue
 
             # Convert this !e conversion to the shell specific env var specifier
-            value = self.shell_formats[self.language]["env_var"].format(field_name)
+            if self.expand == ExpandMode.Remove:
+                value = ""
+            else:
+                value = self.shell_formats[self.language]["env_var"].format(field_name)
             yield (literal_text + value, None, None, None)

hab/parsers/config.py

@@ -3,6 +3,7 @@
 
 from .. import NotSet, utils
 from ..errors import HabError, InvalidAliasError
+from ..formatter import ExpandMode
 from .hab_base import HabBase
 from .meta import hab_property
 
@@ -111,7 +112,7 @@ def launch(self, alias_name, args=None, blocking=False, cls=None, **kwargs):
             env = kwargs["env"]
         else:
             env = dict(os.environ)
-        self.update_environ(env, alias_name)
+        self.update_environ(env, alias_name, expand=ExpandMode.Remove)
         kwargs["env"] = env
 
         # Launch the subprocess using the requested Popen subclass

hab/parsers/hab_base.py

@@ -16,7 +16,7 @@
     InvalidAliasError,
     ReservedVariableNameError,
 )
-from ..formatter import Formatter
+from ..formatter import ExpandMode, Formatter
 from ..site import MergeDict
 from ..solvers import Solver
 from .meta import HabMeta, hab_property
@@ -757,7 +757,14 @@ def uri(self):
             return self._uri
         return self.fullpath
 
-    def update_environ(self, env, alias_name=None, include_global=True, formatter=None):
+    def update_environ(
+        self,
+        env,
+        alias_name=None,
+        include_global=True,
+        formatter=None,
+        expand=ExpandMode.ToShell,
+    ):
         """Updates the given environment variable dictionary to conform with
         the hab environment specification.
 
@@ -772,17 +779,18 @@ def update_environ(self, env, alias_name=None, include_global=True, formatter=No
                 also adds `HAB_FREEZE` if possible.
             formatter (hab.formatter.Formatter, optional): Str formatter class
                 used to format the env var values.
+            ext (str, optional): Defaults to `Platform.default_ext()` if NotSet.
         """
         ext = utils.Platform.default_ext()
         if formatter is None:
             # Make sure to expand environment variables when formatting.
-            formatter = Formatter(ext, expand=True)
+            formatter = Formatter(ext, expand=expand)
 
         def _apply(data):
             for key, value in data.items():
                 if value:
                     value = utils.Platform.collapse_paths(value, ext=ext, key=key)
-                    value = formatter.format(value, key=key, value=value)
+                    value = formatter.format(value, key=key, value=value, hab_env=env)
                     env[key] = value
                 else:
                     env.pop(key, None)

tests/test_formatter.py

@@ -1,7 +1,7 @@
 import pytest
 
 from hab import utils
-from hab.formatter import Formatter
+from hab.formatter import ExpandMode, Formatter
 from hab.parsers import Config
 
 
@@ -37,8 +37,12 @@ def test_env_format(language, expanded, not_expanded, monkeypatch):
     # Check that "!e" is converted to the correct shell specific specifier.
     assert Formatter(language).format(fmt, regular_var="V") == not_expanded
 
-    # Check that "!e" uses the env var value if `expand=True` not the shell specifier.
-    assert Formatter(language, expand=True).format(fmt, regular_var="V") == expanded
+    # Check that "!e" uses the env var value if `expand=ExpandMode.ToShell` not
+    # the shell specifier.
+    assert (
+        Formatter(language, expand=ExpandMode.ToShell).format(fmt, regular_var="V")
+        == expanded
+    )
 
 
 def test_language_from_ext(monkeypatch):