[pypy-svn] pypy default: Support for "pypy somedir" or "pypy file.zip", using the hook

[arigo]

Author: Armin Rigo <arigo at tunes.org>
Branch: 
Changeset: r41928:641b29ed1501
Date: 2011-02-14 20:02 +0100
http://bitbucket.org/pypy/pypy/changeset/641b29ed1501/

Log:	Support for "pypy somedir" or "pypy file.zip", using the hook
	imp._getimporter().

diff --git a/pypy/translator/goal/app_main.py b/pypy/translator/goal/app_main.py
--- a/pypy/translator/goal/app_main.py
+++ b/pypy/translator/goal/app_main.py
@@ -507,10 +507,8 @@
             success = run_toplevel(run_it)
         elif run_module:
             # handle the "-m" command
-            def run_it():
-                import runpy
-                runpy._run_module_as_main(sys.argv[0])
-            success = run_toplevel(run_it)
+            import runpy
+            success = run_toplevel(runpy._run_module_as_main, sys.argv[0])
         elif run_stdin:
             # handle the case where no command/filename/module is specified
             # on the command-line.
@@ -559,15 +557,26 @@
             # assume it's a pyc file only if its name says so.
             # CPython goes to great lengths to detect other cases
             # of pyc file format, but I think it's ok not to care.
+            import imp
             if IS_WINDOWS:
                 filename = filename.lower()
             if filename.endswith('.pyc') or filename.endswith('.pyo'):
-                import imp
-                success = run_toplevel(imp._run_compiled_module, '__main__',
-                                       sys.argv[0], None, mainmodule)
+                args = (imp._run_compiled_module, '__main__',
+                        sys.argv[0], None, mainmodule)
             else:
-                success = run_toplevel(execfile, sys.argv[0],
-                                       mainmodule.__dict__)
+                # maybe it's the name of a directory or a zip file
+                filename = sys.argv[0]
+                importer = imp._getimporter(filename)
+                if not isinstance(importer, imp.NullImporter):
+                    # yes.  put the filename in sys.path[0] and import
+                    # the module __main__
+                    import runpy
+                    sys.path.insert(0, filename)
+                    args = (runpy._run_module_as_main, '__main__')
+                else:
+                    # no.  That's the normal path, "pypy stuff.py".
+                    args = (execfile, filename, mainmodule.__dict__)
+            success = run_toplevel(*args)
 
     except SystemExit, e:
         status = e.code
@@ -637,9 +646,9 @@
         except OSError:
             return None
 
-    # add an emulator for this pypy-only function
+    # add an emulator for these pypy-only or 2.7-only functions
     # (for test_pyc_commandline_argument)
-    import imp
+    import imp, runpy
     def _run_compiled_module(modulename, filename, file, module):
         import os
         assert modulename == '__main__'
@@ -648,7 +657,20 @@
         assert file is None
         assert module.__name__ == '__main__'
         print 'in _run_compiled_module'
+    def _getimporter(path):
+        import os, imp
+        if os.path.isdir(path):
+            return None
+        else:
+            return imp.NullImporter()
+    def _run_module_as_main(module):
+        import sys, os
+        path = os.path.join(sys.path[0], module + '.py')
+        execfile(path, {'__name__': '__main__'})
+
     imp._run_compiled_module = _run_compiled_module
+    imp._getimporter = _getimporter
+    runpy._run_module_as_main = _run_module_as_main
 
     # stick the current sys.path into $PYTHONPATH, so that CPython still
     # finds its own extension modules :-/

diff --git a/pypy/translator/goal/test2/test_app_main.py b/pypy/translator/goal/test2/test_app_main.py
--- a/pypy/translator/goal/test2/test_app_main.py
+++ b/pypy/translator/goal/test2/test_app_main.py
@@ -14,9 +14,9 @@
 app_main = os.path.abspath(app_main)
 
 _counter = 0
-def _get_next_path():
+def _get_next_path(ext='.py'):
     global _counter
-    p = udir.join('demo_test_app_main_%d.py' % (_counter,))
+    p = udir.join('demo_test_app_main_%d%s' % (_counter, ext))
     _counter += 1
     return p
 
@@ -26,7 +26,7 @@
     # return relative path for testing purposes 
     return py.path.local().bestrelpath(p) 
 
-def getscriptpyc(space, source):
+def getscript_pyc(space, source):
     p = _get_next_path()
     p.write(str(py.code.Source(source)))
     w_dir = space.wrap(str(p.dirpath()))
@@ -45,6 +45,13 @@
     assert os.path.isfile(p)   # the .pyc file should have been created above
     return p
 
+def getscript_in_dir(source):
+    pdir = _get_next_path(ext='')
+    p = pdir.ensure(dir=1).join('__main__.py')
+    p.write(str(py.code.Source(source)))
+    # return relative path for testing purposes 
+    return py.path.local().bestrelpath(pdir)
+
 demo_script = getscript("""
     print 'hello'
     print 'Name:', __name__
@@ -699,11 +706,18 @@
         assert data.startswith("[''")
 
     def test_pyc_commandline_argument(self):
-        p = getscriptpyc(self.space, "print 6*7\n")
-        assert p.endswith('.pyc')
+        p = getscript_pyc(self.space, "print 6*7\n")
+        assert os.path.isfile(p) and p.endswith('.pyc')
         data = self.run(p)
         assert data == 'in _run_compiled_module\n'
 
+    def test_main_in_dir_commandline_argument(self):
+        p = getscript_in_dir('print 6*7\n')
+        data = self.run(p)
+        assert data == '42\n'
+        data = self.run(p + os.sep)
+        assert data == '42\n'
+
 
 class AppTestAppMain: