gh-100239: spcialize dict subclasses with no getitem override

python · iritkatriel · Apr 10, 2025 · Apr 10, 2025 · Apr 11, 2025 · Apr 12, 2025
commit bf2d1ddcec26921b52a324a49a567562d135944e
diff --git a/Lib/test/test_opcache.py b/Lib/test/test_opcache.py
@@ -1,3 +1,4 @@
+import collections
 import copy
 import pickle
 import dis
@@ -1729,9 +1730,43 @@ def binary_subscr_dict():
                 self.assertEqual(a[1], 2)
                 self.assertEqual(a[2], 3)
 
-        binary_subscr_dict()
-        self.assert_specialized(binary_subscr_dict, "BINARY_OP_SUBSCR_DICT")
-        self.assert_no_opcode(binary_subscr_dict, "BINARY_OP")
+        def binary_subscr_dict_subclass_defaultdict():
+            for _ in range(_testinternalcapi.SPECIALIZATION_THRESHOLD):
+                a = collections.defaultdict(lambda : 42, {1: 2, 2: 3})
+                self.assertEqual(a[1], 2)
+                self.assertEqual(a[2], 3)
+                self.assertEqual(a[7], 42)
+
+        def binary_subscr_dict_subclass_Counter():
+            for _ in range(_testinternalcapi.SPECIALIZATION_THRESHOLD):
+                a = collections.Counter('abcdeabcdabcaba')
+                self.assertEqual(a['a'], 5)
+                self.assertEqual(a['b'], 4)
+                self.assertEqual(a['m'], 0)
+
+        for f in [binary_subscr_dict,
+                  binary_subscr_dict_subclass_defaultdict,
+                  binary_subscr_dict_subclass_Counter]:
+
+            with self.subTest(f=f):
+                f()
+                self.assert_specialized(f, "BINARY_OP_SUBSCR_DICT")
+                self.assert_no_opcode(f, "BINARY_OP")
+
+        def binary_subscr_dict_subclass_with_subscript_override():
+            class MyDict(dict):
+                def __getitem__(self, key):
+                    return 42
+
+            for _ in range(_testinternalcapi.SPECIALIZATION_THRESHOLD):
+                a = MyDict()
+                self.assertEqual(a['a'], 42)
+                self.assertEqual(a['b'], 42)
+
+        binary_subscr_dict_subclass_with_subscript_override()
+        self.assert_no_opcode(
+            binary_subscr_dict_subclass_with_subscript_override,
+            "BINARY_OP_SUBSCR_DICT")
 
         def binary_subscr_str_int():
             for _ in range(_testinternalcapi.SPECIALIZATION_THRESHOLD):

diff --git a/Objects/typeobject.c b/Objects/typeobject.c
@@ -11184,7 +11184,14 @@ update_one_slot(PyTypeObject *type, pytype_slotdef *p)
         }
         else {
             use_generic = 1;
-            generic = p->function;
+            if (generic == NULL && Py_IS_TYPE(descr, &PyMethodDescr_Type) &&
+                *ptr == ((PyMethodDescrObject *)descr)->d_method->ml_meth)
+            {
+                generic = *ptr;
+            }
+            else {
+                generic = p->function;
+            }
             if (p->function == slot_tp_call) {
                 /* A generic __call__ is incompatible with vectorcall */
                 type_clear_flags(type, Py_TPFLAGS_HAVE_VECTORCALL);

@@ -939,15 +939,14 @@ dummy_func(
             PyObject *sub = PyStackRef_AsPyObjectBorrow(sub_st);
             PyObject *dict = PyStackRef_AsPyObjectBorrow(dict_st);
 
-            DEOPT_IF(!PyDict_CheckExact(dict));
+            DEOPT_IF(!PyDict_Check(dict));
+            DEOPT_IF(!Py_TYPE(dict)->tp_as_mapping);
+            DEOPT_IF(Py_TYPE(dict)->tp_as_mapping->mp_subscript !=
+                     PyDict_Type.tp_as_mapping->mp_subscript);
             STAT_INC(BINARY_OP, hit);
-            PyObject *res_o;
-            int rc = PyDict_GetItemRef(dict, sub, &res_o);
-            if (rc == 0) {
-                _PyErr_SetKeyError(sub);
-            }
+            PyObject *res_o = PyDict_Type.tp_as_mapping->mp_subscript(dict, sub);
             DECREF_INPUTS();
-            ERROR_IF(rc <= 0, error); // not found or error
+            ERROR_IF(res_o == NULL, error); // not found or error
             res = PyStackRef_FromPyObjectSteal(res_o);
         }
 

diff --git a/Python/executor_cases.c.h b/Python/executor_cases.c.h
diff --git a/Python/generated_cases.c.h b/Python/generated_cases.c.h
diff --git a/Python/specialize.c b/Python/specialize.c
@@ -2618,9 +2618,13 @@ _Py_Specialize_BinaryOp(_PyStackRef lhs_st, _PyStackRef rhs_st, _Py_CODEUNIT *in
                     return;
                 }
             }
-            if (PyDict_CheckExact(lhs)) {
-                specialize(instr, BINARY_OP_SUBSCR_DICT);
-                return;
+            if (PyDict_Check(lhs)) {
+                if ((Py_TYPE(lhs)->tp_as_mapping != NULL) &&
+                    (Py_TYPE(lhs)->tp_as_mapping->mp_subscript == PyDict_Type.tp_as_mapping->mp_subscript))
+                {
+                    specialize(instr, BINARY_OP_SUBSCR_DICT);
+                    return;
+                }
             }
             unsigned int tp_version;
             PyTypeObject *container_type = Py_TYPE(lhs);