Skip to content
38 changes: 38 additions & 0 deletions Lib/test/test_zoneinfo/test_zoneinfo.py
Original file line number Diff line number Diff line change
Expand Up @@ -1577,6 +1577,44 @@ class EvilZoneInfo(self.klass):
class CZoneInfoCacheTest(ZoneInfoCacheTest):
module = c_zoneinfo

def test_inconsistent_weak_cache_get(self):
class Cache:
def get(self, key, default=None):
return 1337

class ZI(self.klass):
pass
# Class attribute must be set after class creation
# to override zoneinfo.ZoneInfo.__init_subclass__.
ZI._weak_cache = Cache()

with self.assertRaises(RuntimeError) as te:
ZI("America/Los_Angeles")
self.assertEqual(
str(te.exception),
"Unexpected instance of int in ZI weak cache for key 'America/Los_Angeles'"
)

def test_inconsistent_weak_cache_setdefault(self):
class Cache:
def get(self, key, default=None):
return default
def setdefault(self, key, value):
return 1337

class ZI(self.klass):
pass
# Class attribute must be set after class creation
# to override zoneinfo.ZoneInfo.__init_subclass__.
ZI._weak_cache = Cache()

with self.assertRaises(RuntimeError) as te:
ZI("America/Los_Angeles")
self.assertEqual(
str(te.exception),
"Unexpected instance of int in ZI weak cache for key 'America/Los_Angeles'"
)


class ZoneInfoPickleTest(TzPathUserMixin, ZoneInfoTestBase):
module = py_zoneinfo
Expand Down
Original file line number Diff line number Diff line change
@@ -0,0 +1,2 @@
:mod:`zoneinfo`: fix a crash when instantiating :class:`~zoneinfo.ZoneInfo`
objects for which the internal class-level cache is inconsistent.
11 changes: 10 additions & 1 deletion Modules/_zoneinfo.c
Original file line number Diff line number Diff line change
Expand Up @@ -335,14 +335,23 @@ zoneinfo_ZoneInfo_impl(PyTypeObject *type, PyObject *key)
return NULL;
}

((PyZoneInfo_ZoneInfo *)tmp)->source = SOURCE_CACHE;
instance =
PyObject_CallMethod(weak_cache, "setdefault", "OO", key, tmp);
Py_DECREF(tmp);
if (instance == NULL) {
Py_DECREF(weak_cache);
return NULL;
}
((PyZoneInfo_ZoneInfo *)instance)->source = SOURCE_CACHE;
}

if (!PyObject_TypeCheck(instance, type)) {
PyErr_Format(PyExc_RuntimeError,
"Unexpected instance of %T in %s weak cache for key %R",
instance, _PyType_Name(type), key);
Py_DECREF(instance);
Py_DECREF(weak_cache);
return NULL;
}

update_strong_cache(state, type, key, instance);
Expand Down
Loading