py/objdict: Implement dictionary union (PEP 584).

Implements dictionary union according to PEP 584's specifications, minus
the fact that dictionary entries are not guaranteed to be in insertion
order.  This feature is enabled with MICROPY_CPYTHON_COMPAT.

Includes a new test.

With the assistance of Fangrui Qin <qinf@purdue.edu>

Signed-off-by: Rayane Chatrieux <rayane.chatrieux@gmail.com>
Signed-off-by: Damien George <damien@micropython.org>

3 files changed, 59 insertions, 0 deletions

diff --git a/py/objdict.c b/py/objdict.c
index 7fad5fc8f..baad8a1b9 100644
--- a/py/objdict.c
+++ b/py/objdict.c
@@ -190,6 +190,17 @@ STATIC mp_obj_t dict_binary_op(mp_binary_op_t op, mp_obj_t lhs_in, mp_obj_t rhs_
                 return mp_const_false;
             }
         }
+        #if MICROPY_CPYTHON_COMPAT
+        case MP_BINARY_OP_INPLACE_OR:
+        case MP_BINARY_OP_OR: {
+            if (op == MP_BINARY_OP_OR) {
+                lhs_in = mp_obj_dict_copy(lhs_in);
+            }
+            mp_obj_t dicts[2] = {lhs_in, rhs_in};
+            dict_update(2, dicts, (mp_map_t *)&mp_const_empty_map);
+            return lhs_in;
+        }
+        #endif
         default:
             // op not supported
             return MP_OBJ_NULL;
diff --git a/tests/basics/dict_union.py b/tests/basics/dict_union.py
new file mode 100644
index 000000000..606b1c4c5
--- /dev/null
+++ b/tests/basics/dict_union.py
@@ -0,0 +1,36 @@
+# Test dictionary union.
+#
+# The tests sort the resulting dictionaries for now, since map insertion
+# order is not implemented in MicroPython.
+
+try:
+    {} | {}
+except TypeError:
+    print("SKIP")
+    raise SystemExit
+
+
+def print_sorted_dict(d):
+    print(sorted(d.items()))
+
+
+def test_union(a, b):
+    print_sorted_dict(a | b)
+    print_sorted_dict(b | a)
+    a |= a
+    print_sorted_dict(a)
+    a |= b
+    print_sorted_dict(a)
+
+
+d = {}
+e = {}
+test_union(d, e)
+
+d = {1: "apple"}
+e = {1: "cheese"}
+test_union(d, e)
+
+d = {"spam": 1, "eggs": 2, "cheese": 3}
+e = {"cheese": "cheddar", "aardvark": "Ethel"}
+test_union(d, e)
diff --git a/tests/basics/dict_union.py.exp b/tests/basics/dict_union.py.exp
new file mode 100644
index 000000000..7a8f83bd5
--- /dev/null
+++ b/tests/basics/dict_union.py.exp
@@ -0,0 +1,12 @@
+[]
+[]
+[]
+[]
+[(1, 'cheese')]
+[(1, 'apple')]
+[(1, 'apple')]
+[(1, 'cheese')]
+[('aardvark', 'Ethel'), ('cheese', 'cheddar'), ('eggs', 2), ('spam', 1)]
+[('aardvark', 'Ethel'), ('cheese', 3), ('eggs', 2), ('spam', 1)]
+[('cheese', 3), ('eggs', 2), ('spam', 1)]
+[('aardvark', 'Ethel'), ('cheese', 'cheddar'), ('eggs', 2), ('spam', 1)]