Repository: amoffat/sh
Pull request #466 (2019-03-02 17:50:56)

Fix execution contexts on the REPL
This is a fix for #422 

On the REPL and from the command line, Python doesn't keep source code
in the stack frame for the toplevel.  Therefore we can't inspect it and
look for the variable being assigned to in order to remove it from
sys.modules.

This is a cornercase but we can hack around it by trying to retrieve the
variable name from the code object instead.  This depends more heavily
on Python internals but this was already code which dove into Python's
internals a bit so perhaps this is okay.

<strike>Since I needed to run the tests on Python3.7, this also includes a fix for the test suite on Python3.7</strike>  I pulled the Python-3.7 changes out to their own PR.

abadger: I'm not sure what the issue is trying to build this on pypy.  The test suite passes locally for me on pypy but I have a newer version than is running on travis:
```
$ pypy --version
Python 2.7.13 (1d8462e42c37, Aug 30 2018, 09:58:57)
[PyPy 6.0.0 with GCC 8.2.1 20180801 (Red Hat 8.2.1-2)]
```
https://toshio.fedorapeople.org/sh-passing-tests-on-local-pypy.log
abadger: Okay, I grabbed the travis version of pypy in a VM and I found the problem.  Older versions of pypy don't report parent[1] in the same way as CPython (and newer pypy corrects it).  I'll special case that to get it to work.

Patch:
diff --git a/sh.py b/sh.py
index 4f97a8a..1b3a041 100644
--- a/sh.py
+++ b/sh.py
@@ -3535,12 +3535,30 @@ class SelfWrapper(ModuleType):
         # cached module from sys.modules.  if we don't, it gets re-used, and any
         # old baked params get used, which is not what we want
         parent = inspect.stack()[1]
-        code = parent[4][0].strip()
-        parsed = ast.parse(code)
-        module_name = parsed.body[0].targets[0].id
+        try:
+            code = parent[4][0].strip()
+        except TypeError:
+            # On the REPL or from the commandline, we don't get the source code in the
+            # top stack frame
+            # Older versions of pypy don't set parent[1] the same way as CPython or newer versions
+            # of Pypy so we have to special case that too.
+            if parent[1] in ('<stdin>', '<string>') or (
+                    parent[1] == '<module>' and platform.python_implementation().lower() == 'pypy'):
+                # This depends on things like Python's calling convention and the layout of stack
+                # frames but it's a fix for a bug in a very cornery cornercase so....
+                module_name = parent[0].f_code.co_names[-1]
+        else:
+            parsed = ast.parse(code)
+            try:
+                module_name = parsed.body[0].targets[0].id
+            except Exception:
+                # Diagnose what went wrong
+                if not isinstance(parsed.body[0], ast.Assign):
+                    raise RuntimeError("A new execution context must be assigned to a variable")
+                raise
 
         if module_name == __name__:
-            raise RuntimeError("Cannot use the name 'sh' as an execution context")
+            raise RuntimeError("Cannot use the name '%s' as an execution context" % __name__)
 
         sys.modules.pop(module_name, None)
 
@@ -3572,7 +3590,7 @@ def register_importer():
 
     if not already_registered:
         importer = ModuleImporterFromVariables(
-            restrict_to=["SelfWrapper"],
+            restrict_to=[SelfWrapper.__name__],
         )
         sys.meta_path.insert(0, importer)
 
@@ -3606,17 +3624,9 @@ class ModuleImporterFromVariables(object):
         the rest of this function """
 
         parent_frame = inspect.currentframe().f_back
-
-        while parent_frame and in_importlib(parent_frame):
+        while in_importlib(parent_frame):
             parent_frame = parent_frame.f_back
 
-        # Calling PyImport_ImportModule("some_module"); via the C API may not
-        # have a parent frame. Early-out to avoid in_importlib() trying to
-        # get f_code from None when looking for 'some_module'.
-        # This also happens when using gevent apparently.
-        if not parent_frame:
-            return None
-
         # this line is saying "hey, does mod_fullname exist as a name we've
         # defind previously?"  the purpose of this is to ensure that
         # mod_fullname is really a thing we've defined.  if we haven't defined


Test patch:
diff --git a/test.py b/test.py
index aebb46d..79ca4b1 100644
--- a/test.py
+++ b/test.py
@@ -3092,6 +3092,31 @@ class ExecutionContextTests(unittest.TestCase):
             from _os import path
         self.assertRaises(ImportError, unallowed_import)
 
+    def test_reimport_from_cli(self):
+        # The REPL and CLI both need special handling to create an execution context that is safe to
+        # reimport
+        if IS_PY3:
+            cmdstr = '; '.join(('import sh, io, sys',
+                                'out = io.StringIO()',
+                                '_sh = sh(_out=out)',
+                                'import _sh',
+                                '_sh.echo("-n", "TEST")',
+                                'sys.stderr.write(out.getvalue())',
+                                ))
+        else:
+            cmdstr = '; '.join(('import sh, StringIO, sys',
+                                'out = StringIO.StringIO()',
+                                '_sh = sh(_out=out)',
+                                'import _sh',
+                                '_sh.echo("-n", "TEST")',
+                                'sys.stderr.write(out.getvalue())',
+                                ))
+
+        err = StringIO()
+
+        python('-c', cmdstr, _err=err)
+        self.assertEqual('TEST', err.getvalue())
+
 
 if __name__ == "__main__":
     root = logging.getLogger()