bpo-36144: Implement defaultdict union (GH-18729)

For PEP 585 (this isn't in the PEP but is an obvious follow-up).

1 files changed, 45 insertions, 6 deletions

diff --git a/Modules/_collectionsmodule.c b/Modules/_collectionsmodule.c
index 4d5d874..d0a381d 100644
--- a/Modules/_collectionsmodule.c
+++ b/Modules/_collectionsmodule.c
@@ -1990,6 +1990,13 @@ defdict_missing(defdictobject *dd, PyObject *key)
     return value;
 }
 
+static inline PyObject*
+new_defdict(defdictobject *dd, PyObject *arg)
+{
+    return PyObject_CallFunctionObjArgs((PyObject*)Py_TYPE(dd),
+        dd->default_factory ? dd->default_factory : Py_None, arg, NULL);
+}
+
 PyDoc_STRVAR(defdict_copy_doc, "D.copy() -> a shallow copy of D.");
 
 static PyObject *
@@ -1999,11 +2006,7 @@ defdict_copy(defdictobject *dd, PyObject *Py_UNUSED(ignored))
        whose class constructor has the same signature.  Subclasses that
        define a different constructor signature must override copy().
     */
-
-    if (dd->default_factory == NULL)
-        return PyObject_CallFunctionObjArgs((PyObject*)Py_TYPE(dd), Py_None, dd, NULL);
-    return PyObject_CallFunctionObjArgs((PyObject*)Py_TYPE(dd),
-                                        dd->default_factory, dd, NULL);
+    return new_defdict(dd, (PyObject*)dd);
 }
 
 static PyObject *
@@ -2127,6 +2130,42 @@ defdict_repr(defdictobject *dd)
     return result;
 }
 
+static PyObject*
+defdict_or(PyObject* left, PyObject* right)
+{
+    int left_is_self = PyObject_IsInstance(left, (PyObject*)&defdict_type);
+    if (left_is_self < 0) {
+        return NULL;
+    }
+    PyObject *self, *other;
+    if (left_is_self) {
+        self = left;
+        other = right;
+    }
+    else {
+        self = right;
+        other = left;
+    }
+    if (!PyDict_Check(other)) {
+        Py_RETURN_NOTIMPLEMENTED;
+    }
+    // Like copy(), this calls the object's class.
+    // Override __or__/__ror__ for subclasses with different constructors.
+    PyObject *new = new_defdict((defdictobject*)self, left);
+    if (!new) {
+        return NULL;
+    }
+    if (PyDict_Update(new, right)) {
+        Py_DECREF(new);
+        return NULL;
+    }
+    return new;
+}
+
+static PyNumberMethods defdict_as_number = {
+    .nb_or = defdict_or,
+};
+
 static int
 defdict_traverse(PyObject *self, visitproc visit, void *arg)
 {
@@ -2198,7 +2237,7 @@ static PyTypeObject defdict_type = {
     0,                                  /* tp_setattr */
     0,                                  /* tp_as_async */
     (reprfunc)defdict_repr,             /* tp_repr */
-    0,                                  /* tp_as_number */
+    &defdict_as_number,                 /* tp_as_number */
     0,                                  /* tp_as_sequence */
     0,                                  /* tp_as_mapping */
     0,                                  /* tp_hash */