gh-100239: specialize dict subclasses with no getitem override

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):

Python/bytecodes.c

@@ -971,15 +971,14 @@ dummy_func(
             PyObject *sub = PyStackRef_AsPyObjectBorrow(sub_st);
             PyObject *dict = PyStackRef_AsPyObjectBorrow(dict_st);
 
-            assert(PyDict_CheckExact(dict));
+            assert(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);
         }
 

Python/specialize.c

@@ -2648,9 +2648,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);