Merge pull request #962 from Daetalus/eight-hundred

handle the situation which provide a directory to pyston

src/jit.cpp

@@ -17,6 +17,7 @@
 #include <cstdio>
 #include <cstdlib>
 #include <fcntl.h>
+#include <limits.h>
 #include <readline/history.h>
 #include <readline/readline.h>
 #include <stdint.h>
@@ -265,6 +266,29 @@ static int RunModule(const char* module, int set_argv0) {
     return 0;
 }
 
+static int RunMainFromImporter(const char* filename) {
+    PyObject* argv0 = NULL, * importer = NULL;
+
+    if ((argv0 = PyString_FromString(filename)) && (importer = PyImport_GetImporter(argv0))
+        && (importer->cls != null_importer_cls)) {
+        /* argv0 is usable as an import source, so
+               put it in sys.path[0] and import __main__ */
+        PyObject* sys_path = NULL;
+        if ((sys_path = PySys_GetObject("path")) && !PyList_SetItem(sys_path, 0, argv0)) {
+            Py_INCREF(argv0);
+            Py_DECREF(importer);
+            sys_path = NULL;
+            return RunModule("__main__", 0) != 0;
+        }
+    }
+    Py_XDECREF(argv0);
+    Py_XDECREF(importer);
+    if (PyErr_Occurred()) {
+        PyErr_Print();
+        return 1;
+    }
+    return -1;
+}
 
 static int main(int argc, char** argv) {
     argv0 = argv[0];
@@ -431,9 +455,14 @@ static int main(int argc, char** argv) {
             main_module = createModule(boxString("__main__"), "<string>");
             rtncode = (RunModule(module, 1) != 0);
         } else {
+            main_module = createModule(boxString("__main__"), fn);
             rtncode = 0;
             if (fn != NULL) {
-                llvm::SmallString<128> path;
+                rtncode = RunMainFromImporter(fn);
+            }
+
+            if (rtncode == -1 && fn != NULL) {
+                llvm::SmallString<PATH_MAX> path;
 
                 if (!llvm::sys::fs::exists(fn)) {
                     fprintf(stderr, "[Errno 2] No such file or directory: '%s'\n", fn);
@@ -455,10 +484,10 @@ static int main(int argc, char** argv) {
                 prependToSysPath(real_path);
                 free(real_path);
 
-                main_module = createModule(boxString("__main__"), fn);
                 try {
                     AST_Module* ast = caching_parse_file(fn, /* future_flags = */ 0);
                     compileAndRunModule(ast, main_module);
+                    rtncode = 0;
                 } catch (ExcInfo e) {
                     setCAPIException(e);
                     PyErr_Print();

src/runtime/import.cpp

@@ -35,7 +35,6 @@ namespace pyston {
 static const std::string all_str("__all__");
 static const std::string name_str("__name__");
 static const std::string package_str("__package__");
-static BoxedClass* null_importer_cls;
 
 static void removeModule(BoxedString* name) {
     BoxedDict* d = getSysModulesDict();
@@ -266,6 +265,18 @@ SearchResult findModule(const std::string& name, BoxedString* full_name, BoxedLi
     return SearchResult("", SearchResult::SEARCH_ERROR);
 }
 
+PyObject* PyImport_GetImporter(PyObject* path) noexcept {
+    PyObject* importer = NULL, * path_importer_cache = NULL, * path_hooks = NULL;
+
+    if ((path_importer_cache = PySys_GetObject("path_importer_cache"))) {
+        if ((path_hooks = PySys_GetObject("path_hooks"))) {
+            importer = get_path_importer(path_importer_cache, path_hooks, path);
+        }
+    }
+    Py_XINCREF(importer); /* get_path_importer returns a borrowed reference */
+    return importer;
+}
+
 /* Return the package that an import is being performed in.  If globals comes
    from the module foo.bar.bat (not itself a package), this returns the
    sys.modules entry for foo.bar.  If globals is from a package's __init__.py,

src/runtime/import.h

@@ -19,6 +19,7 @@
 
 namespace pyston {
 
+extern "C" PyObject* PyImport_GetImporter(PyObject* path) noexcept;
 extern "C" Box* import(int level, Box* from_imports, llvm::StringRef module_name);
 extern Box* importModuleLevel(llvm::StringRef module_name, Box* globals, Box* from_imports, int level);
 BoxedModule* importCExtension(BoxedString* full_name, const std::string& last_name, const std::string& path);

src/runtime/types.cpp

@@ -1541,9 +1541,9 @@ void BoxedClosure::gcHandler(GCVisitor* v, Box* b) {
 extern "C" {
 BoxedClass* object_cls, *type_cls, *none_cls, *bool_cls, *int_cls, *float_cls,
     * str_cls = NULL, *function_cls, *instancemethod_cls, *list_cls, *slice_cls, *module_cls, *dict_cls, *tuple_cls,
-      *file_cls, *member_descriptor_cls, *closure_cls, *generator_cls, *complex_cls, *basestring_cls, *property_cls,
-      *staticmethod_cls, *classmethod_cls, *attrwrapper_cls, *pyston_getset_cls, *capi_getset_cls,
-      *builtin_function_or_method_cls, *attrwrapperiter_cls, *set_cls, *frozenset_cls;
+      *file_cls, *member_descriptor_cls, *closure_cls, *generator_cls, *null_importer_cls, *complex_cls,
+      *basestring_cls, *property_cls, *staticmethod_cls, *classmethod_cls, *attrwrapper_cls, *pyston_getset_cls,
+      *capi_getset_cls, *builtin_function_or_method_cls, *attrwrapperiter_cls, *set_cls, *frozenset_cls;
 
 BoxedTuple* EmptyTuple;
 }

src/runtime/types.h

@@ -90,10 +90,10 @@ extern "C" BoxedString* EmptyString;
 extern "C" {
 extern BoxedClass* object_cls, *type_cls, *bool_cls, *int_cls, *long_cls, *float_cls, *str_cls, *function_cls,
     *none_cls, *instancemethod_cls, *list_cls, *slice_cls, *module_cls, *dict_cls, *tuple_cls, *file_cls,
-    *enumerate_cls, *xrange_cls, *member_descriptor_cls, *method_cls, *closure_cls, *generator_cls, *complex_cls,
-    *basestring_cls, *property_cls, *staticmethod_cls, *classmethod_cls, *attrwrapper_cls, *pyston_getset_cls,
-    *capi_getset_cls, *builtin_function_or_method_cls, *set_cls, *frozenset_cls, *code_cls, *frame_cls, *capifunc_cls,
-    *wrapperdescr_cls, *wrapperobject_cls;
+    *enumerate_cls, *xrange_cls, *member_descriptor_cls, *null_importer_cls, *method_cls, *closure_cls, *generator_cls,
+    *complex_cls, *basestring_cls, *property_cls, *staticmethod_cls, *classmethod_cls, *attrwrapper_cls,
+    *pyston_getset_cls, *capi_getset_cls, *builtin_function_or_method_cls, *set_cls, *frozenset_cls, *code_cls,
+    *frame_cls, *capifunc_cls, *wrapperdescr_cls, *wrapperobject_cls;
 }
 #define unicode_cls (&PyUnicode_Type)
 #define memoryview_cls (&PyMemoryView_Type)