Set __classcell__ if it exists

filmor · filmor · commit f3714a27f915 · 2020-05-15T00:44:43.000+02:00
diff --git a/src/runtime/runtime.cs b/src/runtime/runtime.cs
@@ -285,6 +285,11 @@ internal static void Initialize(bool initSigs = false)
                 () => PyFloatType = IntPtr.Zero);
             XDecref(op);
 
+            op = PyCell_New();
+            SetPyMember(ref PyCellType, PyObject_Type(op),
+                () => PyCellType = IntPtr.Zero);
+            // XDecref(op);
+
 #if !PYTHON2
             PyClassType = IntPtr.Zero;
             PyInstanceType = IntPtr.Zero;
@@ -452,6 +457,7 @@ private static void ResetPyMembers()
         internal static IntPtr PyBoolType;
         internal static IntPtr PyNoneType;
         internal static IntPtr PyTypeType;
+        internal static IntPtr PyCellType;
 
         internal static IntPtr Py_NoSiteFlag;
 
@@ -1965,6 +1971,18 @@ internal static IntPtr PyMem_Realloc(IntPtr ptr, long size)
         [DllImport(_PythonDll, CallingConvention = CallingConvention.Cdecl)]
         internal static extern void PyErr_Print();
 
+        //====================================================================
+        // Cell API
+        //====================================================================
+
+        [DllImport(_PythonDll, CallingConvention = CallingConvention.Cdecl)]
+        internal static extern IntPtr PyCell_New();
+
+        [DllImport(_PythonDll, CallingConvention = CallingConvention.Cdecl)]
+        internal static extern NewReference PyCell_Get(IntPtr cell);
+
+        [DllImport(_PythonDll, CallingConvention = CallingConvention.Cdecl)]
+        internal static extern int PyCell_Set(IntPtr cell, IntPtr value);
 
         //====================================================================
         // Miscellaneous
diff --git a/src/runtime/typemanager.cs b/src/runtime/typemanager.cs
@@ -280,6 +280,21 @@ internal static IntPtr CreateSubType(IntPtr py_name, IntPtr py_base_type, IntPtr
                 IntPtr cls_dict = Marshal.ReadIntPtr(py_type, TypeOffset.tp_dict);
                 Runtime.PyDict_Update(cls_dict, py_dict);
 
+                // Update the __classcell__ if it exists
+                IntPtr classCell = Runtime.PyString_FromString("__classcell__");
+                if (Runtime.PySequence_Contains(cls_dict, classCell) > 0)
+                {
+                    IntPtr cell = Runtime.PyDict_GetItem(cls_dict, classCell);
+                    if (Runtime.PyObject_TypeCheck(cell, Runtime.PyCellType))
+                    {
+                        Runtime.XIncref(py_type);
+                        Runtime.PyCell_Set(cell, py_type);
+                    }
+                    Runtime.PyDict_DelItem(cls_dict, classCell);
+                }
+
+                Runtime.XDecref(classCell);
+
                 return py_type;
             }
             catch (Exception e)
