feat(abstractions): enhance HeadersCollection usability and dict interoperability

packages/abstractions/kiota_abstractions/headers_collection.py

@@ -1,14 +1,20 @@
 from __future__ import annotations
 
-from typing import Union
+from collections.abc import Iterable, Iterator, Mapping
+from typing import Optional, Union
 
 
 class HeadersCollection():
     "Represents a collection of request/response headers"
     SINGLE_VALUE_HEADERS: set[str] = {"content-type", "content-encoding", "content-length"}
 
-    def __init__(self) -> None:
+    def __init__(
+        self,
+        headers: Optional[Union[HeadersCollection, Mapping[str, Union[str, Iterable[str]]]]] = None,
+    ) -> None:
         self._headers: dict[str, set[str]] = {}
+        if headers is not None:
+            self.add_all(headers)
 
     def try_get(self, key: str) -> Union[bool, set[str]]:
         """Gets the header values corresponding to a specific header name.
@@ -72,17 +78,30 @@
             return True
         return False
 
-    def add_all(self, headers: HeadersCollection) -> None:
+    def add_all(
+        self,
+        headers: Union[HeadersCollection, Mapping[str, Union[str, Iterable[str]]]],
+    ) -> None:
         """Adds the specified headers to the collection.
 
         Args:
-            headers (dict[str, str]): The headers to add.
+            headers: The headers to add.
         """
-        if not headers:
+        if headers is None:
             raise ValueError("Headers cannot be null")
-        for key, values in headers.get_all().items():
-            for value in values:
-                self.add(key, value)
+        if isinstance(headers, HeadersCollection):
+            for key, coll_values in headers.get_all().items():
+                self.add(key, list(coll_values))
+        elif isinstance(headers, Mapping):
+            for key, map_val in headers.items():
+                if isinstance(map_val, str):
+                    self.add(key, map_val)
+                elif isinstance(map_val, Iterable):
+                    self.add(key, [str(item) for item in map_val])
+                else:
+                    self.add(key, str(map_val))
+        else:
+            raise TypeError("Headers must be a HeadersCollection or Mapping")
 
     def add(self, header_name: str, header_values: Union[str, list[str]]) -> None:
         """Adds values to the header with the specified name.
@@ -177,3 +196,46 @@
             raise ValueError("Header name cannot be null")
         key = key.lower()
         return key in self._headers
+
+    def __contains__(self, key: object) -> bool:
+        if not isinstance(key, str) or not key:
+            return False
+        return self.contains(key)
+
+    def __getitem__(self, key: str) -> set[str]:
+        if not self.contains(key):
+            raise KeyError(key)
+        return self.get(key)
+
+    def __setitem__(self, key: str, value: Union[str, Iterable[str]]) -> None:
+        if not key:
+            raise ValueError("Header name cannot be null")
+        if value is None:
+            raise ValueError("Header values cannot be null")
+        header_name = key.lower()
+        if isinstance(value, str):
+            header_values = [value]
+        elif isinstance(value, Iterable):
+            header_values = list(value)
+        else:
+            header_values = [str(value)]
+
+        if not header_values:
+            self._headers.pop(header_name, None)
+            return
+
+        if header_name in self.SINGLE_VALUE_HEADERS:
+            self._headers[header_name] = {header_values[0]}
+        else:
+            self._headers[header_name] = set(header_values)
+
+    def __delitem__(self, key: str) -> None:
+        if not self.contains(key):
+            raise KeyError(key)
+        self.remove(key)
+
+    def __len__(self) -> int:
+        return self.count()
+
+    def __iter__(self) -> Iterator[str]:
+        return iter(self._headers)

packages/abstractions/tests/test_request_header.py

@@ -181,6 +181,8 @@ def test_adds_headers_from_instance():
     assert headers.contains("header4")
     assert headers.count() == 4
     assert headers.keys() == ["header1", "header2", "header3", "header4"]
+
+
 def test_add_all_replaces_single_value_header():
     """Tests that add_all keeps a single value header to one value
     """
@@ -190,3 +192,114 @@ def test_add_all_replaces_single_value_header():
     other.add("Content-Type", "text/plain")
     headers.add_all(other)
     assert headers.get("content-type") == {"text/plain"}
+
+
+def test_init_with_dict():
+    """Tests initialization of HeadersCollection with a dict"""
+    headers = HeadersCollection({"header1": "value1", "header2": ["value2", "value3"]})
+    assert headers.count() == 2
+    assert "header1" in headers
+    assert "header2" in headers
+    assert headers["header1"] == {"value1"}
+    assert headers["header2"] == {"value2", "value3"}
+
+
+def test_init_with_headers_collection():
+    """Tests initialization of HeadersCollection with another HeadersCollection"""
+    source = HeadersCollection({"header1": "value1"})
+    copy_headers = HeadersCollection(source)
+    assert copy_headers.count() == 1
+    assert "header1" in copy_headers
+    assert copy_headers["header1"] == {"value1"}
+
+
+def test_init_with_empty_inputs():
+    """Tests that initializing with empty mapping or empty HeadersCollection does not fail"""
+    headers_dict = HeadersCollection({})
+    assert headers_dict.count() == 0
+    assert len(headers_dict) == 0
+
+    empty_source = HeadersCollection()
+    headers_copy = HeadersCollection(empty_source)
+    assert headers_copy.count() == 0
+    assert len(headers_copy) == 0
+
+
+def test_init_iterable_single_value_header():
+    """Tests that initializing with an iterable for single-value headers retains only the first value"""
+    headers = HeadersCollection({"Content-Length": ["10", "20"]})
+    assert headers.get("content-length") == {"10"}
+    assert headers["content-length"] == {"10"}
+
+
+def test_add_all_with_dict():
+    """Tests adding headers from a dictionary"""
+    headers = HeadersCollection()
+    headers.add_all({"header1": "value1", "header2": ["value2", "value3"]})
+    assert headers.count() == 2
+    assert headers.get("header1") == {"value1"}
+    assert headers.get("header2") == {"value2", "value3"}
+
+
+def test_add_all_empty_inputs():
+    """Tests that add_all with empty collection or empty mapping is a no-op and does not raise"""
+    headers = HeadersCollection({"header1": "value1"})
+    headers.add_all({})
+    assert headers.count() == 1
+    headers.add_all(HeadersCollection())
+    assert headers.count() == 1
+
+
+def test_container_methods():
+    """Tests container magic methods (__contains__, __getitem__, __setitem__, __len__, __iter__)"""
+    headers = HeadersCollection()
+    headers["header1"] = "value1"
+    headers["header2"] = ["val2", "val3"]
+    assert len(headers) == 2
+    assert "header1" in headers
+    assert "header2" in headers
+    assert "nonexistent" not in headers
+    assert headers["header1"] == {"value1"}
+    assert headers["header2"] == {"val2", "val3"}
+    assert list(headers) == ["header1", "header2"]
+    with pytest.raises(KeyError):
+        _ = headers["nonexistent"]
+
+
+def test_subscript_assignment_replaces_values():
+    """Tests that subscript assignment replaces existing values rather than appending, across casing"""
+    headers = HeadersCollection()
+    headers["Authorization"] = "old"
+    assert headers["authorization"] == {"old"}
+    headers["authorization"] = "new"
+    assert headers["authorization"] == {"new"}
+    assert headers.get("authorization") == {"new"}
+    assert len(headers) == 1
+
+    headers["AUTHORIZATION"] = ["token1", "token2"]
+    assert headers["authorization"] == {"token1", "token2"}
+    assert len(headers) == 1
+
+
+def test_subscript_assignment_validation():
+    """Tests that invalid key or value raises ValueError before modifying stored headers"""
+    headers = HeadersCollection()
+    headers["Authorization"] = "valid"
+
+    with pytest.raises(ValueError):
+        headers[None] = "value"
+    with pytest.raises(ValueError):
+        headers[""] = "value"
+    with pytest.raises(ValueError):
+        headers["Authorization"] = None
+
+    # Stored value remains intact
+    assert headers["authorization"] == {"valid"}
+
+
+def test_subscript_assignment_single_value_header():
+    """Tests that subscript assignment respects SINGLE_VALUE_HEADERS"""
+    headers = HeadersCollection()
+    headers["Content-Type"] = ["application/json", "text/plain"]
+    assert headers["content-type"] == {"application/json"}
+

packages/abstractions/tests/test_request_information.py

@@ -26,7 +26,7 @@ def test_initialization():
     assert not request_info.url_template
     assert not request_info.http_method
     assert not request_info.content
-    assert request_info.headers
+    assert request_info.headers is not None
     assert request_info.RAW_URL_KEY == 'request-raw-url'
     assert request_info.BINARY_CONTENT_TYPE == 'application/octet-stream'
     assert request_info.CONTENT_TYPE_HEADER == 'Content-Type'