From b0926115659768f0ae08fb98d351796d82afa049 Mon Sep 17 00:00:00 2001
From: Hristo Venev <hristo@venev.name>
Date: Thu, 28 Jan 2021 20:53:35 +0200
Subject: [PATCH] Don't downgrade unicode-only dicts to mixed on non-unicode
 lookups

---
 Objects/dictobject.c | 11 ++++++++---
 1 file changed, 8 insertions(+), 3 deletions(-)

diff --git a/Objects/dictobject.c b/Objects/dictobject.c
index faee6bc901..96a1bf5e9b 100644
--- a/Objects/dictobject.c
+++ b/Objects/dictobject.c
@@ -438,6 +438,12 @@ dictkeys_set_index(PyDictKeysObject *keys, Py_ssize_t i, Py_ssize_t ix)
         (d)->ma_keys->dk_lookup = lookdict_unicode; \
     }
 
+#define MAYBE_NONUNICODE(d, key) \
+    if (!PyUnicode_CheckExact(key) && (d)->ma_keys->dk_lookup != lookdict) { \
+        assert((d)->ma_keys->dk_lookup == lookdict_unicode || (d)->ma_keys->dk_lookup == lookdict_unicode_nodummy); \
+        (d)->ma_keys->dk_lookup = lookdict; \
+    }
+
 /* This immutable, empty PyDictKeysObject is used for PyDict_Clear()
  * (which cannot fail and thus can do no allocation).
  */
@@ -835,7 +841,6 @@ lookdict_unicode(PyDictObject *mp, PyObject *key,
        unicodes is to override __eq__, and for speed we don't cater to
        that here. */
     if (!PyUnicode_CheckExact(key)) {
-        mp->ma_keys->dk_lookup = lookdict;
         return lookdict(mp, key, hash, value_addr);
     }
 
@@ -878,7 +883,6 @@ lookdict_unicode_nodummy(PyDictObject *mp, PyObject *key,
        unicodes is to override __eq__, and for speed we don't cater to
        that here. */
     if (!PyUnicode_CheckExact(key)) {
-        mp->ma_keys->dk_lookup = lookdict;
         return lookdict(mp, key, hash, value_addr);
     }
 
@@ -1062,7 +1066,6 @@ insertdict(PyDictObject *mp, PyObject *key, Py_hash_t hash, PyObject *value)
     if (ix == DKIX_ERROR)
         goto Fail;
 
-    assert(PyUnicode_CheckExact(key) || mp->ma_keys->dk_lookup == lookdict);
     MAINTAIN_TRACKING(mp, key, value);
 
     /* When insertion order is different from shared key, we can't share
@@ -1084,6 +1087,7 @@ insertdict(PyDictObject *mp, PyObject *key, Py_hash_t hash, PyObject *value)
             if (insertion_resize(mp) < 0)
                 goto Fail;
         }
+        MAYBE_NONUNICODE(mp, key);
         Py_ssize_t hashpos = find_empty_slot(mp->ma_keys, hash);
         ep = &DK_ENTRIES(mp->ma_keys)[mp->ma_keys->dk_nentries];
         dictkeys_set_index(mp->ma_keys, hashpos, mp->ma_keys->dk_nentries);
@@ -2965,6 +2969,7 @@ PyDict_SetDefault(PyObject *d, PyObject *key, PyObject *defaultobj)
                 return NULL;
             }
         }
+        MAYBE_NONUNICODE(mp, key);
         Py_ssize_t hashpos = find_empty_slot(mp->ma_keys, hash);
         ep0 = DK_ENTRIES(mp->ma_keys);
         ep = &ep0[mp->ma_keys->dk_nentries];
-- 
2.29.2