unix/main: Implement -m option for packages.

pfalcon · pfalcon · commit edc02bd952ed · 2017-05-09T14:22:21.000+03:00
diff --git a/py/builtinimport.c b/py/builtinimport.c
@@ -429,8 +429,13 @@ mp_obj_t mp_builtin___import__(size_t n_args, const mp_obj_t *args) {
 
                 // if args[3] (fromtuple) has magic value False, set up
                 // this module for command-line "-m" option (set module's
-                // name to __main__ instead of real name).
-                if (i == mod_len && fromtuple == mp_const_false) {
+                // name to __main__ instead of real name). Do this only
+                // for *modules* however - packages never have their names
+                // replaced, instead they're -m'ed using a special __main__
+                // submodule in them. (This all apparently is done to not
+                // touch package name itself, which is important for future
+                // imports).
+                if (i == mod_len && fromtuple == mp_const_false && stat != MP_IMPORT_STAT_DIR) {
                     mp_obj_module_t *o = MP_OBJ_TO_PTR(module_obj);
                     mp_obj_dict_store(MP_OBJ_FROM_PTR(o->globals), MP_OBJ_NEW_QSTR(MP_QSTR___name__), MP_OBJ_NEW_QSTR(MP_QSTR___main__));
                     #if MICROPY_CPYTHON_COMPAT
diff --git a/unix/main.c b/unix/main.c
@@ -553,6 +553,9 @@ MP_NOINLINE int main_(int argc, char **argv) {
 
                 mp_obj_t mod;
                 nlr_buf_t nlr;
+                bool subpkg_tried = false;
+
+            reimport:
                 if (nlr_push(&nlr) == 0) {
                     mod = mp_builtin___import__(MP_ARRAY_SIZE(import_args), import_args);
                     nlr_pop();
@@ -561,11 +564,17 @@ MP_NOINLINE int main_(int argc, char **argv) {
                     return handle_uncaught_exception(nlr.ret_val) & 0xff;
                 }
 
-                if (mp_obj_is_package(mod)) {
-                    // TODO
-                    mp_printf(&mp_stderr_print, "%s: -m for packages not yet implemented\n", argv[0]);
-                    exit(1);
+                if (mp_obj_is_package(mod) && !subpkg_tried) {
+                    subpkg_tried = true;
+                    vstr_t vstr;
+                    int len = strlen(argv[a + 1]);
+                    vstr_init(&vstr, len + sizeof(".__main__"));
+                    vstr_add_strn(&vstr, argv[a + 1], len);
+                    vstr_add_strn(&vstr, ".__main__", sizeof(".__main__") - 1);
+                    import_args[0] = mp_obj_new_str_from_vstr(&mp_type_str, &vstr);
+                    goto reimport;
                 }
+
                 ret = 0;
                 break;
             } else if (strcmp(argv[a], "-X") == 0) {
