Add union operators to MutableHeaders (#1239)

starlette/datastructures.py

@@ -611,6 +611,19 @@ def __delitem__(self, key: str) -> None:
         for idx in reversed(pop_indexes):
             del self._list[idx]
 
+    def __ior__(self, other: dict) -> "MutableHeaders":
+        if not isinstance(other, typing.Mapping):
+            raise NotImplementedError
+        self.update(other)
+        return self
+
+    def __or__(self, other: dict) -> "MutableHeaders":
+        if not isinstance(other, typing.Mapping):
+            raise NotImplementedError
+        new = self.mutablecopy()
+        new.update(other)
+        return new
+
     @property
     def raw(self) -> typing.List[typing.Tuple[bytes, bytes]]:
         return self._list

tests/test_datastructures.py

@@ -162,6 +162,48 @@ def test_mutable_headers():
     assert h.raw == [(b"b", b"4")]
 
 
+def test_mutable_headers_merge():
+    h = MutableHeaders()
+    h = h | MutableHeaders({"a": "1"})
+    assert isinstance(h, MutableHeaders)
+    assert dict(h) == {"a": "1"}
+    assert h.items() == [("a", "1")]
+    assert h.raw == [(b"a", b"1")]
+
+    with pytest.raises(NotImplementedError):
+        h |= {"not_mapping"}
+
+    with pytest.raises(NotImplementedError):
+        h | {"not_mapping"}
+
+
+def test_mutable_headers_merge_dict():
+    h = MutableHeaders()
+    h = h | {"a": "1"}
+    assert isinstance(h, MutableHeaders)
+    assert dict(h) == {"a": "1"}
+    assert h.items() == [("a", "1")]
+    assert h.raw == [(b"a", b"1")]
+
+
+def test_mutable_headers_update():
+    h = MutableHeaders()
+    h |= MutableHeaders({"a": "1"})
+    assert isinstance(h, MutableHeaders)
+    assert dict(h) == {"a": "1"}
+    assert h.items() == [("a", "1")]
+    assert h.raw == [(b"a", b"1")]
+
+
+def test_mutable_headers_update_dict():
+    h = MutableHeaders()
+    h |= {"a": "1"}
+    assert isinstance(h, MutableHeaders)
+    assert dict(h) == {"a": "1"}
+    assert h.items() == [("a", "1")]
+    assert h.raw == [(b"a", b"1")]
+
+
 def test_headers_mutablecopy():
     h = Headers(raw=[(b"a", b"123"), (b"a", b"456"), (b"b", b"789")])
     c = h.mutablecopy()