[PYTHON] Enable environment scoping (apache#33)

tqchen · Jul 12, 2018 · 9f0e8ff · 9f0e8ff
1 parent 7ab574b
commit 9f0e8ff
Show file tree

Hide file tree

Showing 2 changed files with 18 additions and 0 deletions.
diff --git a/vta/python/vta/environment.py b/vta/python/vta/environment.py
@@ -143,6 +143,15 @@ def __init__(self, cfg):
  self.mock_mode = False
  self._mock_env = None
  self._dev_ctx = None
+ self._last_env = None
+
+ def __enter__(self):
+ self._last_env = Environment.current
+ Environment.current = self
+ return self
+
+ def __exit__(self, ptype, value, trace):
+ Environment.current = self._last_env
 
  def pkg_config(self):
  """PkgConfig instance"""

diff --git a/vta/tests/python/unittest/test_environment.py b/vta/tests/python/unittest/test_environment.py
@@ -6,6 +6,15 @@ def test_env():
  mock = env.mock
  assert mock.alu == "skip_alu"
 
+def test_env_scope():
+ env = vta.get_env()
+ cfg = env.pkg_config().cfg_dict
+ cfg["TARGET"] = "xyz"
+ with vta.Environment(cfg):
+ assert vta.get_env().TARGET == "xyz"
+ assert vta.get_env().TARGET == env.TARGET
+
 
 if __name__ == "__main__":
  test_env()
+ test_env_scope()