From 06e0963e3020ccf9b1b63ea874d0894326c393ff Mon Sep 17 00:00:00 2001
From: manlix <manlix@yandex.ru>
Date: Tue, 6 Jul 2021 18:51:43 +0300
Subject: [PATCH] Add union operators to MutableHeaders (#1239)

---
 starlette/datastructures.py  | 13 +++++++++++
 tests/test_datastructures.py | 42 ++++++++++++++++++++++++++++++++++++
 2 files changed, 55 insertions(+)

diff --git a/starlette/datastructures.py b/starlette/datastructures.py
index 64f964a915..fc5e7ef056 100644
--- a/starlette/datastructures.py
+++ b/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
diff --git a/tests/test_datastructures.py b/tests/test_datastructures.py
index bb71ba870c..7af16804a0 100644
--- a/tests/test_datastructures.py
+++ b/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()