Split from PR 958: restoring the __import__ after shutdown. (#993)

benoithudson · filmor · commit 5f2e2e2f81d0 · 2019-11-21T10:10:58.000+01:00
When C# shuts down we should restore Python to its original state.
diff --git a/src/runtime/importhook.cs b/src/runtime/importhook.cs
@@ -26,25 +26,64 @@ internal static void InitializeModuleDef()
 #endif
 
         /// <summary>
-        /// Initialization performed on startup of the Python runtime.
+        /// Get a <i>New reference</i> to the builtins module.
         /// </summary>
-        internal static void Initialize()
+        static IntPtr GetNewRefToBuiltins()
         {
-            // Initialize the Python <--> CLR module hook. We replace the
-            // built-in Python __import__ with our own. This isn't ideal,
-            // but it provides the most "Pythonic" way of dealing with CLR
-            // modules (Python doesn't provide a way to emulate packages).
-            IntPtr dict = Runtime.PyImport_GetModuleDict();
+            if (Runtime.IsPython3)
+            {
+                return Runtime.PyImport_ImportModule("builtins");
+            }
+            else
+            {
+                // dict is a borrowed ref, no need to decref
+                IntPtr dict = Runtime.PyImport_GetModuleDict();
 
-            IntPtr mod = Runtime.IsPython3
-                ? Runtime.PyImport_ImportModule("builtins")
-                : Runtime.PyDict_GetItemString(dict, "__builtin__");
+                // GetItemString is a borrowed ref; incref to get a new ref
+                IntPtr builtins = Runtime.PyDict_GetItemString(dict, "__builtin__");
+                Runtime.XIncref(builtins);
+                return builtins;
+            }
+        }
 
-            py_import = Runtime.PyObject_GetAttrString(mod, "__import__");
+        /// <summary>
+        /// Initialize just the __import__ hook itself.
+        /// </summary>
+        static void InitImport()
+        {
+            // We replace the built-in Python __import__ with our own: first
+            // look in CLR modules, then if we don't find any call the default
+            // Python __import__.
+            IntPtr builtins = GetNewRefToBuiltins();
+            py_import = Runtime.PyObject_GetAttrString(builtins, "__import__");
             hook = new MethodWrapper(typeof(ImportHook), "__import__", "TernaryFunc");
-            Runtime.PyObject_SetAttrString(mod, "__import__", hook.ptr);
+            Runtime.PyObject_SetAttrString(builtins, "__import__", hook.ptr);
             Runtime.XDecref(hook.ptr);
+            Runtime.XDecref(builtins);
+        }
+
+        /// <summary>
+        /// Restore the __import__ hook.
+        /// </summary>
+        static void RestoreImport()
+        {
+            IntPtr builtins = GetNewRefToBuiltins();
+
+            Runtime.PyObject_SetAttrString(builtins, "__import__", py_import);
+            Runtime.XDecref(py_import);
+            py_import = IntPtr.Zero;
+
+            Runtime.XDecref(builtins);
+        }
+
+        /// <summary>
+        /// Initialization performed on startup of the Python runtime.
+        /// </summary>
+        internal static void Initialize()
+        {
+            InitImport();
 
+            // Initialize the clr module and tell Python about it.
             root = new CLRModule();
 
 #if PYTHON3
@@ -62,6 +101,7 @@ internal static void Initialize()
             Runtime.XIncref(root.pyHandle); // we are using the module two times
             py_clr_module = root.pyHandle; // Alias handle for PY2/PY3
 #endif
+            IntPtr dict = Runtime.PyImport_GetModuleDict();
             Runtime.PyDict_SetItemString(dict, "CLR", py_clr_module);
             Runtime.PyDict_SetItemString(dict, "clr", py_clr_module);
         }
@@ -74,9 +114,10 @@ internal static void Shutdown()
         {
             if (Runtime.Py_IsInitialized() != 0)
             {
+                RestoreImport();
+
                 Runtime.XDecref(py_clr_module);
                 Runtime.XDecref(root.pyHandle);
-                Runtime.XDecref(py_import);
             }
         }
 
