From f3e4c505d1a2d29caf0a6095445c53f0b493208f Mon Sep 17 00:00:00 2001
From: Rayane Chatrieux <rayane.chatrieux@gmail.com>
Date: Mon, 11 Apr 2022 17:35:56 -0400
Subject: [PATCH] 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>
---
 py/objdict.c                   | 11 +++++++++++
 tests/basics/dict_union.py     | 36 ++++++++++++++++++++++++++++++++++
 tests/basics/dict_union.py.exp | 12 ++++++++++++
 3 files changed, 59 insertions(+)
 create mode 100644 tests/basics/dict_union.py
 create mode 100644 tests/basics/dict_union.py.exp

diff --git a/py/objdict.c b/py/objdict.c
index 7fad5fc8f4..baad8a1b9e 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 0000000000..606b1c4c52
--- /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 0000000000..7a8f83bd54
--- /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)]