1 files changed, 88 insertions, 9 deletions

diff --git a/intern/mantaflow/intern/MANTA_main.cpp b/intern/mantaflow/intern/MANTA_main.cpp
index 282fbdb3f77..fc14c909f4d 100644
--- a/intern/mantaflow/intern/MANTA_main.cpp
+++ b/intern/mantaflow/intern/MANTA_main.cpp
@@ -562,15 +562,66 @@ MANTA::~MANTA()
   pythonCommands.push_back(finalString);
   result = runPythonString(pythonCommands);
 
+  /* WARNING: this causes crash on exit in the `cycles_volume_cpu/smoke_color` test,
+   * freeing a single modifier ends up clearing the shared module.
+   * For this to be handled properly there would need to be a initialize/free
+   * function for global data. */
+#if 0
+  MANTA::terminateMantaflow();
+#endif
+
   BLI_assert(result);
   UNUSED_VARS(result);
 }
 
 /**
- * Store a pointer to the __main__ module used by mantaflow. This is necessary, because sometimes
- * Blender will overwrite that module. That happens when e.g. scripts are executed in the text
- * editor.
- *
+ * Copied from `PyC_DefaultNameSpace` in Blender.
+ * with some differences:
+ * - Doesn't touch `sys.modules`, use #manta_python_main_module_activate instead.
+ * - Returns the module instead of the modules `dict`.
+ * */
+static PyObject *manta_python_main_module_create(const char *filename)
+{
+  PyObject *builtins = PyEval_GetBuiltins();
+  PyObject *mod_main = PyModule_New("__main__");
+  PyModule_AddStringConstant(mod_main, "__name__", "__main__");
+  if (filename) {
+    /* __file__ mainly for nice UI'ness
+     * NOTE: this won't map to a real file when executing text-blocks and buttons. */
+    PyModule_AddObject(mod_main, "__file__", PyUnicode_InternFromString(filename));
+  }
+  PyModule_AddObject(mod_main, "__builtins__", builtins);
+  Py_INCREF(builtins); /* AddObject steals a reference */
+  return mod_main;
+}
+
+static void manta_python_main_module_activate(PyObject *mod_main)
+{
+  PyObject *modules = PyImport_GetModuleDict();
+  PyObject *main_mod_cmp = PyDict_GetItemString(modules, "__main__");
+  if (mod_main == main_mod_cmp) {
+    return;
+  }
+  /* NOTE: we could remove the reference to `mod_main` here, but as it's know to be removed
+   * accept that there is temporarily an extra reference. */
+  PyDict_SetItemString(modules, "__main__", mod_main);
+}
+
+static void manta_python_main_module_backup(PyObject **r_main_mod)
+{
+  PyObject *modules = PyImport_GetModuleDict();
+  *r_main_mod = PyDict_GetItemString(modules, "__main__");
+  Py_XINCREF(*r_main_mod); /* don't free */
+}
+
+static void manta_python_main_module_restore(PyObject *main_mod)
+{
+  PyObject *modules = PyImport_GetModuleDict();
+  PyDict_SetItemString(modules, "__main__", main_mod);
+  Py_XDECREF(main_mod);
+}
+
+/**
  * Mantaflow stores many variables in the globals() dict of the __main__ module. To be able to
  * access these variables, the same __main__ module has to be used every time.
  *
@@ -579,14 +630,35 @@ MANTA::~MANTA()
  */
 static PyObject *manta_main_module = nullptr;
 
+static void manta_python_main_module_clear()
+{
+  if (manta_main_module) {
+    Py_DECREF(manta_main_module);
+    manta_main_module = nullptr;
+  }
+}
+
+static PyObject *manta_python_main_module_ensure()
+{
+  if (!manta_main_module) {
+    manta_main_module = manta_python_main_module_create("<manta_namespace>");
+  }
+  return manta_main_module;
+}
+
 bool MANTA::runPythonString(vector<string> commands)
 {
   bool success = true;
   PyGILState_STATE gilstate = PyGILState_Ensure();
 
-  if (manta_main_module == nullptr) {
-    manta_main_module = PyImport_ImportModule("__main__");
-  }
+  /* Temporarily set `sys.modules["__main__"]` as some Python modules expect this. */
+  PyObject *main_mod_backup;
+  manta_python_main_module_backup(&main_mod_backup);
+
+  /* If we never want to run this when the module isn't initialize,
+   * assign with `manta_python_main_module_ensure()`. */
+  BLI_assert(manta_main_module != nullptr);
+  manta_python_main_module_activate(manta_main_module);
 
   for (vector<string>::iterator it = commands.begin(); it != commands.end(); ++it) {
     string command = *it;
@@ -605,6 +677,9 @@ bool MANTA::runPythonString(vector<string> commands)
       Py_DECREF(return_value);
     }
   }
+
+  manta_python_main_module_restore(main_mod_backup);
+
   PyGILState_Release(gilstate);
 
   BLI_assert(success);
@@ -622,7 +697,10 @@ void MANTA::initializeMantaflow()
   /* Initialize extension classes and wrappers. */
   srand(0);
   PyGILState_STATE gilstate = PyGILState_Ensure();
-  Pb::setup(filename, fill); /* Namespace from Mantaflow (registry). */
+
+  PyObject *manta_main_module = manta_python_main_module_ensure();
+  PyObject *globals_dict = PyModule_GetDict(manta_main_module);
+  Pb::setup(false, filename, fill, globals_dict); /* Namespace from Mantaflow (registry). */
   PyGILState_Release(gilstate);
 }
 
@@ -632,7 +710,8 @@ void MANTA::terminateMantaflow()
     cout << "Fluid: Releasing Mantaflow framework" << endl;
 
   PyGILState_STATE gilstate = PyGILState_Ensure();
-  Pb::finalize(); /* Namespace from Mantaflow (registry). */
+  Pb::finalize(false); /* Namespace from Mantaflow (registry). */
+  manta_python_main_module_clear();
   PyGILState_Release(gilstate);
 }