Nicer __main__ support for Python 3

ISSUE: DEVTOOLS-4841

ref:d402badab719bf5a7bb3aef4e470723ba1a709aa

Author: mityada

build/plugins/pybuild.py

@@ -215,22 +215,19 @@ def onpy_srcs(unit, *args):
         # Unsupported but legal PROTO_LIBRARY arguments.
         elif arg == 'GLOBAL' or arg.endswith('.gztproto'):
             pass
-        # XXX https://st.yandex-team.ru/DEVTOOLS-4841
-        elif py3 and (arg == '__main__.py' or arg.endswith('/__main__.py')):
-            unit.onfix_python_main([arg, '__real_main__.py'])
-            unit.onpy_srcs(['TOP_LEVEL', '__real_main__.py'])
-            unit.onpy_main(['__real_main__:real_main_func'])
         # Sources.
         else:
             main_mod = arg == 'MAIN'
             if main_mod:
                 arg = next(args)
 
             if '=' in arg:
+                main_py = False
                 path, mod = arg.split('=', 1)
             else:
                 path = arg
-                if arg == '__main__.py' or arg.endswith('/__main__.py'):
+                main_py = (path == '__main__.py' or path.endswith('/__main__.py'))
+                if not py3 and main_py:
                     mod = '__main__'
                 else:
                     if arg.startswith('../'):
@@ -241,7 +238,9 @@ def onpy_srcs(unit, *args):
                     mod = ns + stripext(arg).replace('/', '.')
 
             if main_mod:
-                unit.onpy_main(mod)
+                py_main(unit, mod + ":main")
+            elif py3 and main_py:
+                py_main(unit, mod)
 
             pathmod = (path, mod)
 
@@ -455,6 +454,11 @@ def onpy_register(unit, *args):
             py_register(unit, name, py3)
 
 
+def py_main(unit, arg):
+    py_program(unit, is_py3(unit))
+    unit.onresource(['-', 'PY_MAIN={}'.format(arg)])
+
+
 def onpy_main(unit, arg):
     """
         @usage: PY_MAIN(pkg.mod[:func])
@@ -466,5 +470,4 @@ def onpy_main(unit, arg):
     if ':' not in arg:
         arg += ':main'
 
-    py_program(unit, is_py3(unit))
-    unit.onresource(['-', 'PY_MAIN={}'.format(arg)])
+    py_main(unit, arg)

build/scripts/fix_python_main.py

@@ -1994,10 +1994,6 @@ macro COPY_FILE(File, Destination) {
     .CMD=$COPY_CMD ${input:File} ${output;noauto:Destination} ${kv;hide:"p CP"} ${kv;hide:"pc light-cyan"}
 }
 
-macro FIX_PYTHON_MAIN(File, Destination) {
-    .CMD=$YMAKE_PYTHON ${input:"build/scripts/fix_python_main.py"} ${input:File} ${output;noauto:Destination} ${kv;hide:"p CP"} ${kv;hide:"pc light-cyan"}
-}
-
 macro BUNDLE_TARGET(Target, Destination) {
     .CMD=$MOVE_FILE ${result:Target} ${output:Destination} ${kv;hide:"p BN"} ${kv;hide:"pc light-cyan"}
 }

build/ymake.core.conf

@@ -19,6 +19,55 @@ void unsetenv(const char* name) {
 }
 #endif
 
+static int RunModule(const char *modname)
+{
+    PyObject *module, *runpy, *runmodule, *runargs, *result;
+    runpy = PyImport_ImportModule("runpy");
+    if (runpy == NULL) {
+        fprintf(stderr, "Could not import runpy module\n");
+        PyErr_Print();
+        return -1;
+    }
+    runmodule = PyObject_GetAttrString(runpy, "_run_module_as_main");
+    if (runmodule == NULL) {
+        fprintf(stderr, "Could not access runpy._run_module_as_main\n");
+        PyErr_Print();
+        Py_DECREF(runpy);
+        return -1;
+    }
+    module = PyUnicode_FromString(modname);
+    if (module == NULL) {
+        fprintf(stderr, "Could not convert module name to unicode\n");
+        PyErr_Print();
+        Py_DECREF(runpy);
+        Py_DECREF(runmodule);
+        return -1;
+    }
+    runargs = Py_BuildValue("(Oi)", module, 0);
+    if (runargs == NULL) {
+        fprintf(stderr,
+            "Could not create arguments for runpy._run_module_as_main\n");
+        PyErr_Print();
+        Py_DECREF(runpy);
+        Py_DECREF(runmodule);
+        Py_DECREF(module);
+        return -1;
+    }
+    result = PyObject_Call(runmodule, runargs, NULL);
+    if (result == NULL) {
+        PyErr_Print();
+    }
+    Py_DECREF(runpy);
+    Py_DECREF(runmodule);
+    Py_DECREF(module);
+    Py_DECREF(runargs);
+    if (result == NULL) {
+        return -1;
+    }
+    Py_DECREF(result);
+    return 0;
+}
+
 int main(int argc, char** argv) {
     int i, sts = 1;
     char* oldloc = NULL;
@@ -111,12 +160,14 @@ int main(int argc, char** argv) {
         func_name = colon + 1;
     }
 
-    PyObject* module = PyImport_ImportModule(module_name);
-
-    if (module == NULL) {
-        PyErr_Print();
+    if (!func_name) {
+        sts = RunModule(module_name);
     } else {
-        if (func_name) {
+        PyObject* module = PyImport_ImportModule(module_name);
+
+        if (module == NULL) {
+            PyErr_Print();
+        } else {
             PyObject* value = PyObject_CallMethod(module, func_name, NULL);
 
             if (value == NULL) {
@@ -125,11 +176,9 @@ int main(int argc, char** argv) {
                 Py_DECREF(value);
                 sts = 0;
             }
-        } else {
-            sts = 0;
-        }
 
-        Py_DECREF(module);
+            Py_DECREF(module);
+        }
     }
 
     if (Py_FinalizeEx() < 0) {