gh-91052: Add C API for watching dictionaries

carljm · carljm · commit 04472325877f · 2022-10-03T16:10:06.000-07:00
diff --git a/Doc/c-api/dict.rst b/Doc/c-api/dict.rst
@@ -238,3 +238,44 @@ Dictionary Objects
           for key, value in seq2:
               if override or key not in a:
                   a[key] = value
+
+.. c:function:: int PyDict_AddWatcher(PyDict_WatchCallback callback)
+
+   Register *callback* as a dictionary watcher. Return a non-negative integer
+   id which must be passed to future calls to :c:func:`PyDict_Watch`. In case
+   of error (e.g. no more watcher IDs available), return ``-1`` and set an
+   exception.
+
+.. c:function:: int PyDict_Watch(int watcher_id, PyObject *dict)
+
+   Mark dictionary *dict* as watched. The callback granted *watcher_id* by
+   :c:func:`PyDict_AddWatcher` will be called when *dict* is modified or
+   deallocated.
+
+.. c:type:: PyDict_WatchEvent
+
+   Enumeration of possible dictionary watcher events: ``PyDict_EVENT_ADDED``,
+   ``PyDict_EVENT_MODIFIED``, ``PyDict_EVENT_DELETED``, ``PyDict_EVENT_CLONED``,
+   ``PyDict_EVENT_CLEARED``, or ``PyDict_EVENT_DEALLOCED``.
+
+.. c:type:: void (*PyDict_WatchCallback)(PyDict_WatchEvent event, PyObject *dict, PyObject *key, PyObject *new_value)
+
+   Type of a dict watcher callback function.
+
+   If *event* is ``PyDict_EVENT_CLEARED`` or ``PyDict_EVENT_DEALLOCED``, both
+   *key* and *new_value* will be ``NULL``. If *event* is
+   ``PyDict_EVENT_ADDED`` or ``PyDict_EVENT_MODIFIED``, *new_value* will be the
+   new value for *key*. If *event* is ``PyDict_EVENT_DELETED``, *key* is being
+   deleted from the dictionary and *new_value* will be ``NULL``.
+
+   ``PyDict_EVENT_CLONED`` occurs when *dict* was previously empty and another
+   dict is merged into it. To maintain efficiency of this operation, per-key
+   ``PyDict_EVENT_ADDED`` events are not issued in this case; instead a
+   single ``PyDict_EVENT_CLONED`` is issued, and *key* will be the source
+   dictionary.
+
+   The callback may inspect but should not modify *dict*; doing so could have
+   unpredictable effects, including infinite recursion.
+
+   Callbacks occur before the notified modification to *dict* takes place, so
+   the prior state of *dict* can be inspected.
diff --git a/Include/cpython/dictobject.h b/Include/cpython/dictobject.h
@@ -83,3 +83,25 @@ typedef struct {
 
 PyAPI_FUNC(PyObject *) _PyDictView_New(PyObject *, PyTypeObject *);
 PyAPI_FUNC(PyObject *) _PyDictView_Intersect(PyObject* self, PyObject *other);
+
+/* Dictionary watchers */
+
+typedef enum {
+    PyDict_EVENT_ADDED,
+    PyDict_EVENT_MODIFIED,
+    PyDict_EVENT_DELETED,
+    PyDict_EVENT_CLONED,
+    PyDict_EVENT_CLEARED,
+    PyDict_EVENT_DEALLOCED,
+} PyDict_WatchEvent;
+
+// Callback to be invoked when a watched dict is cleared, dealloced, or modified.
+// In clear/dealloc case, key and new_value will be NULL. Otherwise, new_value will be the
+// new value for key, NULL if key is being deleted.
+typedef void(*PyDict_WatchCallback)(PyDict_WatchEvent event, PyObject* dict, PyObject* key, PyObject* new_value);
+
+// Register a dict-watcher callback
+PyAPI_FUNC(int) PyDict_AddWatcher(PyDict_WatchCallback callback);
+
+// Mark given dictionary as "watched" (callback will be called if it is modified)
+PyAPI_FUNC(int) PyDict_Watch(int watcher_id, PyObject* dict);
diff --git a/Include/internal/pycore_dict.h b/Include/internal/pycore_dict.h
@@ -154,7 +154,32 @@ struct _dictvalues {
 
 extern uint64_t _pydict_global_version;
 
-#define DICT_NEXT_VERSION() (++_pydict_global_version)
+#define DICT_MAX_WATCHERS 8
+#define DICT_VERSION_MASK 255
+#define DICT_VERSION_INCREMENT 256
+
+#define DICT_NEXT_VERSION() (_pydict_global_version += DICT_VERSION_INCREMENT)
+
+void
+_PyDict_SendEvent(int watcher_bits,
+                  PyDict_WatchEvent event,
+                  PyDictObject *mp,
+                  PyObject *key,
+                  PyObject *value);
+
+static inline uint64_t
+_PyDict_NotifyEvent(PyDict_WatchEvent event,
+                    PyDictObject *mp,
+                    PyObject *key,
+                    PyObject *value)
+{
+    int watcher_bits = mp->ma_version_tag & DICT_VERSION_MASK;
+    if (watcher_bits) {
+        _PyDict_SendEvent(watcher_bits, event, mp, key, value);
+        return DICT_NEXT_VERSION() | watcher_bits;
+    }
+    return DICT_NEXT_VERSION();
+}
 
 extern PyObject *_PyObject_MakeDictFromInstanceAttributes(PyObject *obj, PyDictValues *values);
 extern PyObject *_PyDict_FromItems(
diff --git a/Include/internal/pycore_interp.h b/Include/internal/pycore_interp.h
@@ -144,6 +144,8 @@ struct _is {
     // Initialized to _PyEval_EvalFrameDefault().
     _PyFrameEvalFunction eval_frame;
 
+    void *dict_watchers[8];
+
     Py_ssize_t co_extra_user_count;
     freefunc co_extra_freefuncs[MAX_CO_EXTRA_USERS];
 
diff --git a/Misc/NEWS.d/next/Core and Builtins/2022-10-03-16-12-39.gh-issue-91052.MsYL9d.rst b/Misc/NEWS.d/next/Core and Builtins/2022-10-03-16-12-39.gh-issue-91052.MsYL9d.rst
@@ -0,0 +1 @@
+Add API for subscribing to modification events on selected dictionaries.
diff --git a/Modules/_testcapimodule.c b/Modules/_testcapimodule.c
@@ -5169,6 +5169,206 @@ test_tstate_capi(PyObject *self, PyObject *Py_UNUSED(args))
 }
 
 
+// Test dict watching
+static PyObject *g_dict_watch_events;
+
+static void
+dict_watch_callback(PyDict_WatchEvent event,
+                    PyObject *dict,
+                    PyObject *key,
+                    PyObject *new_value)
+{
+    PyObject *msg;
+    switch(event) {
+        case PyDict_EVENT_CLEARED:
+            msg = PyUnicode_FromString("clear");
+            break;
+        case PyDict_EVENT_DEALLOCED:
+            msg = PyUnicode_FromString("dealloc");
+            break;
+        case PyDict_EVENT_CLONED:
+            msg = PyUnicode_FromString("clone");
+            break;
+        case PyDict_EVENT_ADDED:
+            msg = PyUnicode_FromFormat("new:%S:%S", key, new_value);
+            break;
+        case PyDict_EVENT_MODIFIED:
+            msg = PyUnicode_FromFormat("mod:%S:%S", key, new_value);
+            break;
+        case PyDict_EVENT_DELETED:
+            msg = PyUnicode_FromFormat("del:%S", key);
+            break;
+        default:
+            msg = PyUnicode_FromString("unknown");
+    }
+    assert(PyList_Check(g_dict_watch_events));
+    PyList_Append(g_dict_watch_events, msg);
+}
+
+static int
+dict_watch_assert(Py_ssize_t expected_num_events,
+                  const char *expected_last_msg)
+{
+    char buf[512];
+    Py_ssize_t actual_num_events = PyList_Size(g_dict_watch_events);
+    if (expected_num_events != actual_num_events) {
+        snprintf(buf,
+                 512,
+                 "got %d dict watch events, expected %d",
+                 (int)actual_num_events,
+                 (int)expected_num_events);
+        raiseTestError("test_watch_dict", (const char *)&buf);
+        return -1;
+    }
+    PyObject *last_msg = PyList_GetItem(g_dict_watch_events,
+                                        PyList_Size(g_dict_watch_events)-1);
+    if (PyUnicode_CompareWithASCIIString(last_msg, expected_last_msg)) {
+        snprintf(buf,
+                 512,
+                 "last event is '%s', expected '%s'",
+                 PyUnicode_AsUTF8(last_msg),
+                 expected_last_msg);
+        raiseTestError("test_watch_dict", (const char *)&buf);
+        return -1;
+    }
+    return 0;
+}
+
+static int
+try_watch(int watcher_id, PyObject *obj) {
+    if (PyDict_Watch(watcher_id, obj)) {
+        raiseTestError("test_watch_dict", "PyDict_Watch() failed on dict");
+        return -1;
+    }
+    return 0;
+}
+
+static int
+dict_watch_assert_error(int watcher_id, PyObject *obj, const char *fail_msg)
+{
+    if (!PyDict_Watch(watcher_id, obj)) {
+        raiseTestError("test_watch_dict", fail_msg);
+        return -1;
+    } else if (!PyErr_Occurred()) {
+        raiseTestError("test_watch_dict", "PyDict_Watch() returned error code without exception set");
+        return -1;
+    } else {
+        PyErr_Clear();
+    }
+    return 0;
+}
+
+static PyObject *
+test_watch_dict(PyObject *self, PyObject *Py_UNUSED(args))
+{
+    PyObject *watched = PyDict_New();
+    PyObject *unwatched = PyDict_New();
+    PyObject *one = PyLong_FromLong(1);
+    PyObject *two = PyLong_FromLong(2);
+    PyObject *key1 = PyUnicode_FromString("key1");
+    PyObject *key2 = PyUnicode_FromString("key2");
+
+    g_dict_watch_events = PyList_New(0);
+
+    int wid = PyDict_AddWatcher(dict_watch_callback);
+    if (try_watch(wid, watched)) {
+        return NULL;
+    }
+
+    PyDict_SetItem(unwatched, key1, two);
+    PyDict_Merge(watched, unwatched, 1);
+
+    if (dict_watch_assert(1, "clone")) {
+        return NULL;
+    }
+
+    PyDict_SetItem(watched, key1, one);
+    PyDict_SetItem(unwatched, key1, one);
+
+    if (dict_watch_assert(2, "mod:key1:1")) {
+        return NULL;
+    }
+
+    PyDict_SetItemString(watched, "key1", two);
+    PyDict_SetItemString(unwatched, "key1", two);
+
+    if (dict_watch_assert(3, "mod:key1:2")) {
+        return NULL;
+    }
+
+    PyDict_SetItem(watched, key2, one);
+    PyDict_SetItem(unwatched, key2, one);
+
+    if (dict_watch_assert(4, "new:key2:1")) {
+        return NULL;
+    }
+
+    _PyDict_Pop(watched, key2, Py_None);
+    _PyDict_Pop(unwatched, key2, Py_None);
+
+    if (dict_watch_assert(5, "del:key2")) {
+        return NULL;
+    }
+
+    PyDict_DelItemString(watched, "key1");
+    PyDict_DelItemString(unwatched, "key1");
+
+    if (dict_watch_assert(6, "del:key1")) {
+        return NULL;
+    }
+
+    PyDict_SetDefault(watched, key1, one);
+    PyDict_SetDefault(unwatched, key1, one);
+
+    if (dict_watch_assert(7, "new:key1:1")) {
+        return NULL;
+    }
+
+    PyDict_Clear(watched);
+    PyDict_Clear(unwatched);
+
+    if (dict_watch_assert(8, "clear")) {
+        return NULL;
+    }
+
+    PyObject *copy = PyDict_Copy(watched);
+    // copied dict is not watched, so this does not add an event
+    Py_CLEAR(copy);
+
+    Py_CLEAR(watched);
+
+    if (dict_watch_assert(9, "dealloc")) {
+        return NULL;
+    }
+
+    // it is an error to try to watch a non-dict
+    if (dict_watch_assert_error(wid, one, "PyDict_Watch() succeeded on non-dict")) {
+        return NULL;
+    }
+
+    // It is an error to pass an out-of-range watcher ID
+    if (dict_watch_assert_error(-1, unwatched, "PyDict_Watch() succeeded on negative watcher ID")) { 
+        return NULL;
+    }
+    if (dict_watch_assert_error(8, unwatched, "PyDict_Watch() succeeded on too-large watcher ID")) {
+        return NULL;
+    }
+
+    // It is an error to pass a never-registered watcher ID
+    if (dict_watch_assert_error(7, unwatched, "PyDict_Watch() succeeded on unused watcher ID")) { 
+        return NULL;
+    }
+
+    Py_CLEAR(unwatched);
+    Py_CLEAR(g_dict_watch_events);
+    Py_DECREF(one);
+    Py_DECREF(two);
+    Py_DECREF(key1);
+    Py_DECREF(key2);
+    Py_RETURN_NONE;
+}
+
+
 // Test PyFloat_Pack2(), PyFloat_Pack4() and PyFloat_Pack8()
 static PyObject *
 test_float_pack(PyObject *self, PyObject *args)
@@ -5762,6 +5962,7 @@ static PyMethodDef TestMethods[] = {
     {"settrace_to_record", settrace_to_record, METH_O, NULL},
     {"test_macros", test_macros, METH_NOARGS, NULL},
     {"clear_managed_dict", clear_managed_dict, METH_O, NULL},
+    {"test_watch_dict", test_watch_dict, METH_NOARGS, NULL},
     {NULL, NULL} /* sentinel */
 };
 
diff --git a/Objects/dictobject.c b/Objects/dictobject.c
diff --git a/Python/ceval.c b/Python/ceval.c

Original file line number	Diff line number	Diff line change
`@@ -0,0 +1 @@`
	`1`	`+Add API for subscribing to modification events on selected dictionaries.`