From ca1a768e420579c70998455fe82e89f42a813807 Mon Sep 17 00:00:00 2001
From: Alex Mykyta <amykyta3@users.noreply.github.com>
Date: Wed, 25 Oct 2023 21:01:42 -0700
Subject: [PATCH] Add support for env variables in .f files. #35

---
 src/peakrdl/main.py         | 20 +++++++++++++++++++-
 test/test_basics.py         |  1 +
 test/testdata/dump_nested.f |  2 +-
 3 files changed, 21 insertions(+), 2 deletions(-)

diff --git a/src/peakrdl/main.py b/src/peakrdl/main.py
index e9253c6..a2eab17 100644
--- a/src/peakrdl/main.py
+++ b/src/peakrdl/main.py
@@ -2,7 +2,8 @@
 import sys
 import os
 import shlex
-from typing import TYPE_CHECKING, List, Dict, Optional, Set
+import re
+from typing import TYPE_CHECKING, List, Dict, Optional, Set, Match
 
 from systemrdl import RDLCompileError
 
@@ -89,9 +90,26 @@ def expand_file_args(argv: List[str], _pathlist: Optional[Set[str]] = None) -> L
     return new_argv
 
 
+def expand_arg_vars(argv: List[str]) -> List[str]:
+    """
+    Expand environment variables in args
+    """
+    pattern = re.compile(r"\$(\w+|\{[^}]*\})")
+    def repl(m: Match) -> str:
+        k = m.group(1)
+        if k.startswith("{") and k.endswith("}"):
+            k = k[1:-1]
+
+        v = os.environ.get(k, m.group(0))
+        return v
+
+    return [pattern.sub(repl, arg) for arg in argv]
+
+
 def main() -> None:
     # manually expand any -f argfiles first
     argv = expand_file_args(sys.argv[1:])
+    argv = expand_arg_vars(argv)
 
     cfg = load_cfg(argv)
 
diff --git a/test/test_basics.py b/test/test_basics.py
index 75d8413..d98d488 100644
--- a/test/test_basics.py
+++ b/test/test_basics.py
@@ -61,6 +61,7 @@ def test_input_dne(self):
         ], expects_error=True)
 
     def test_f_argfile(self):
+        os.environ["PEAKRDL_TOP_TEST_ENVVAR"] = "nested"
         self.run_commandline([
             '-f', os.path.join(self.testdata_dir, "dump_nested.f"),
         ])
diff --git a/test/testdata/dump_nested.f b/test/testdata/dump_nested.f
index 9b5ae09..f81ffee 100644
--- a/test/testdata/dump_nested.f
+++ b/test/testdata/dump_nested.f
@@ -1 +1 @@
-dump testdata/parameters.rdl --top nested
+dump testdata/parameters.rdl --top ${PEAKRDL_TOP_TEST_ENVVAR}