Starlark: Add dict union operators | and |=

Implements the Starlark spec addition of bazelbuild/starlark#215.

Work towards #12457

Closes #14540.

PiperOrigin-RevId: 431737269

Author: fmeum

src/main/java/net/starlark/java/eval/Dict.java

@@ -72,7 +72,7 @@
             + "0 in d, \"a\" in d  # (True, False)\n"
             + "[(k, v) for k, v in d.items()]  # [(0, \"a\"), (1, 1), (2, \"b\")]\n"
             + "</pre>\n"
-            + "<p>There are three ways to construct a dictionary:\n"
+            + "<p>There are four ways to construct a dictionary:\n"
             + "<ol>\n"
             + "<li>A dictionary expression <code>{k: v, ...}</code> yields a new dictionary with"
             + " the specified key/value entries, inserted in the order they appear in the"
@@ -90,6 +90,17 @@
             + " a dictionary containing the specified entries, which are inserted in argument"
             + " order, positional arguments before named. As with comprehensions, duplicate keys"
             + " are permitted.\n"
+            + "<li>The union expression <code>x | y</code> yields a new dictionary by combining two"
+            + " existing dictionaries. If the two dictionaries have a key <code>k</code> in common,"
+            + " the right hand side dictionary's value of the key (in other words,"
+            + " <code>y[k]</code>) wins. The <code>|=</code> variant of the union operator modifies"
+            + " a dictionary in-place. Example:<br>"
+            + "<pre class=language-python>"
+            + "d = {\"foo\": \"FOO\", \"bar\": \"BAR\"} | {\"foo\": \"FOO2\", \"baz\": \"BAZ\"}\n"
+            + "# d == {\"foo\": \"FOO2\", \"bar\": \"BAR\", \"baz\": \"BAZ\"}\n"
+            + "d = {\"a\": 1, \"b\": 2}\n"
+            + "d |= {\"b\": 3, \"c\": 4}\n"
+            + "# d == {\"a\": 1, \"b\": 3, \"c\": 4}</pre>"
             + "</ol>")
 public class Dict<K, V>
     implements Map<K, V>,

src/main/java/net/starlark/java/eval/Eval.java

@@ -458,6 +458,14 @@ private static Object inplaceBinaryOp(StarlarkThread.Frame fr, TokenKind op, Obj
       StarlarkList<?> list = (StarlarkList) x;
       list.extend(y);
       return list;
+    } else if (op == TokenKind.PIPE && x instanceof Dict && y instanceof Dict) {
+      // dict |= dict merges the contents of the second dict into the first.
+      @SuppressWarnings("unchecked")
+      Dict<Object, Object> xDict = (Dict<Object, Object>) x;
+      @SuppressWarnings("unchecked")
+      Dict<Object, Object> yDict = (Dict<Object, Object>) y;
+      xDict.putEntries(yDict);
+      return xDict;
     }
     return EvalUtils.binaryOp(op, x, y, fr.thread);
   }

src/main/java/net/starlark/java/eval/EvalUtils.java

@@ -126,6 +126,11 @@ static Object binaryOp(TokenKind op, Object x, Object y, StarlarkThread starlark
             // int | int
             return StarlarkInt.or((StarlarkInt) x, (StarlarkInt) y);
           }
+        } else if (x instanceof Dict) {
+          if (y instanceof Dict) {
+            // dict | dict
+            return Dict.builder().putAll((Dict<?, ?>) x).putAll((Dict<?, ?>) y).build(mu);
+          }
         }
         break;
 

src/test/java/net/starlark/java/eval/testdata/dict.star

@@ -75,6 +75,63 @@ d.update((["f", 0],), f = 6)
 expected = {"a": 1, "b": 2, "c": 3, "d": 4, "e": 5, "f": 6}
 assert_eq(d, expected)
 
+# dict union operator |
+
+empty_dict = dict()
+dict_with_a_b = dict(a=1, b=[1, 2])
+dict_with_b = dict(b=[1, 2])
+dict_with_other_b = dict(b=[3, 4])
+
+assert_eq(empty_dict | dict_with_a_b, dict_with_a_b)
+# Verify iteration order.
+assert_eq((empty_dict | dict_with_a_b).items(), dict_with_a_b.items())
+assert_eq(dict_with_a_b | empty_dict, dict_with_a_b)
+assert_eq((dict_with_a_b | empty_dict).items(), dict_with_a_b.items())
+assert_eq(dict_with_b | dict_with_a_b, dict_with_a_b)
+assert_eq((dict_with_b | dict_with_a_b).items(), dict(b=[1, 2], a=1).items())
+assert_eq(dict_with_a_b | dict_with_b, dict_with_a_b)
+assert_eq((dict_with_a_b | dict_with_b).items(), dict_with_a_b.items())
+assert_eq(dict_with_b | dict_with_other_b, dict_with_other_b)
+assert_eq((dict_with_b | dict_with_other_b).items(), dict_with_other_b.items())
+assert_eq(dict_with_other_b | dict_with_b, dict_with_b)
+assert_eq((dict_with_other_b | dict_with_b).items(), dict_with_b.items())
+
+assert_eq(empty_dict, dict())
+assert_eq(dict_with_b, dict(b=[1,2]))
+
+assert_fails(lambda: dict() | [], "unsupported binary operation")
+
+# dict union assignment operator |=
+
+def test_dict_union_assignment():
+    empty_dict = dict()
+    empty_dict |= {"a": 1}
+    empty_dict |= {"b": 2}
+    empty_dict |= {"c": "3", 7: 4}
+    empty_dict |= {"b": "5", "e": 6}
+    expected_1 = {"a": 1, "b": "5", "c": "3", 7: 4, "e": 6}
+    assert_eq(empty_dict, expected_1)
+    assert_eq(empty_dict.items(), expected_1.items())
+
+    dict_a = {8: 1, "b": 2}
+    dict_b = {"b": 1, "c": 6}
+    dict_c = {"d": 7}
+    dict_tuple = {(5, "a"): ("c", 8)}
+    dict_a |= dict_b
+    dict_c |= dict_a
+    dict_c |= dict_tuple
+    expected_2 = {"d": 7, 8: 1, "b": 1, "c": 6, (5, "a"): ("c", 8)}
+    assert_eq(dict_c, expected_2)
+    assert_eq(dict_c.items(), expected_2.items())
+    assert_eq(dict_b, {"b": 1, "c": 6})
+
+test_dict_union_assignment()
+
+def dict_union_assignment_type_mismatch():
+    some_dict = dict()
+    some_dict |= []
+
+assert_fails(dict_union_assignment_type_mismatch, "unsupported binary operation")
 
 # creation with repeated keys