Merge pull request #23 from leonardt/fix-env

Fix env
leonardt · Nov 13, 2019 · a9f0104 · a9f0104
2 parents bd1d97c + f3d954c
commit a9f0104
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 3 deletions.
diff --git a/ast_tools/stack.py b/ast_tools/stack.py
@@ -102,7 +102,10 @@ def inspect_enclosing_env(
     def wrapped_0(*args, **kwargs):
         exec(_SKIP_FRAME_DEBUG_STMT)
 
-        _st = st or get_symbol_table(list(itertools.chain(decorators, [wrapped_0])))
+        _st = get_symbol_table(list(itertools.chain(decorators, [wrapped_0])))
+        if st is not None:
+            _st.locals.update(st)
+
         env = dict(_st.globals)
         env.update(_st.locals)
         return fn(env, *args, **kwargs)

diff --git a/tests/test_stack.py b/tests/test_stack.py
@@ -56,9 +56,11 @@ def test(env):
     test()
 
 def test_custom_env():
+    MAGIC1 = 'foo'
     def test(env):
-        assert env['MAGIC'] == 'bar'
+        assert env['MAGIC1'] == 'foo'
+        assert env['MAGIC2'] == 'bar'
 
-    st = stack.SymbolTable(locals={},globals={'MAGIC':'bar'})
+    st = stack.SymbolTable(locals={},globals={'MAGIC2':'bar'})
     test = stack.inspect_enclosing_env(test, st=st)
     test()