Broaden header inputs to mapping types across clients

Co-authored-by: Shri Sukhani <shrisukhani@users.noreply.github.com>

Author: cursoragent

hyperbrowser/client/async_client.py

@@ -1,4 +1,4 @@
-from typing import Dict, Optional
+from typing import Mapping, Optional
 
 from ..config import ClientConfig
 from ..transport.async_transport import AsyncTransport
@@ -23,7 +23,7 @@ def __init__(
         config: Optional[ClientConfig] = None,
         api_key: Optional[str] = None,
         base_url: Optional[str] = None,
-        headers: Optional[Dict[str, str]] = None,
+        headers: Optional[Mapping[str, str]] = None,
         timeout: Optional[int] = 30,
     ):
         super().__init__(AsyncTransport, config, api_key, base_url, headers)

hyperbrowser/client/base.py

@@ -1,5 +1,5 @@
 import os
-from typing import Dict, Optional, Type, Union
+from typing import Mapping, Optional, Type, Union
 
 from hyperbrowser.exceptions import HyperbrowserError
 from ..config import ClientConfig
@@ -15,7 +15,7 @@ def __init__(
         config: Optional[ClientConfig] = None,
         api_key: Optional[str] = None,
         base_url: Optional[str] = None,
-        headers: Optional[Dict[str, str]] = None,
+        headers: Optional[Mapping[str, str]] = None,
     ):
         if config is not None and any(
             value is not None for value in (api_key, base_url, headers)

hyperbrowser/client/sync.py

@@ -1,4 +1,4 @@
-from typing import Dict, Optional
+from typing import Mapping, Optional
 
 from ..config import ClientConfig
 from ..transport.sync import SyncTransport
@@ -23,7 +23,7 @@ def __init__(
         config: Optional[ClientConfig] = None,
         api_key: Optional[str] = None,
         base_url: Optional[str] = None,
-        headers: Optional[Dict[str, str]] = None,
+        headers: Optional[Mapping[str, str]] = None,
         timeout: Optional[int] = 30,
     ):
         super().__init__(SyncTransport, config, api_key, base_url, headers)

hyperbrowser/config.py

@@ -12,7 +12,7 @@ class ClientConfig:
 
     api_key: str
     base_url: str = "https://api.hyperbrowser.ai"
-    headers: Optional[Dict[str, str]] = None
+    headers: Optional[Mapping[str, str]] = None
 
     def __post_init__(self) -> None:
         if not isinstance(self.api_key, str):

hyperbrowser/transport/async_transport.py

@@ -1,6 +1,6 @@
 import json
 import httpx
-from typing import Optional
+from typing import Mapping, Optional
 
 from hyperbrowser.exceptions import HyperbrowserError
 from hyperbrowser.version import __version__
@@ -10,7 +10,7 @@
 class AsyncTransport(AsyncTransportStrategy):
     """Asynchronous transport implementation using httpx"""
 
-    def __init__(self, api_key: str, headers: Optional[dict] = None):
+    def __init__(self, api_key: str, headers: Optional[Mapping[str, str]] = None):
         merged_headers = {
             "x-api-key": api_key,
             "User-Agent": f"hyperbrowser-python-sdk/{__version__}",

hyperbrowser/transport/base.py

@@ -1,5 +1,5 @@
 from abc import ABC, abstractmethod
-from typing import Generic, Optional, Type, TypeVar, Union
+from typing import Generic, Mapping, Optional, Type, TypeVar, Union
 
 from hyperbrowser.exceptions import HyperbrowserError
 
@@ -37,7 +37,7 @@ class SyncTransportStrategy(ABC):
     """Abstract base class for synchronous transport implementations"""
 
     @abstractmethod
-    def __init__(self, api_key: str, headers: Optional[dict] = None): ...
+    def __init__(self, api_key: str, headers: Optional[Mapping[str, str]] = None): ...
 
     @abstractmethod
     def close(self) -> None: ...
@@ -63,7 +63,7 @@ class AsyncTransportStrategy(ABC):
     """Abstract base class for asynchronous transport implementations"""
 
     @abstractmethod
-    def __init__(self, api_key: str, headers: Optional[dict] = None): ...
+    def __init__(self, api_key: str, headers: Optional[Mapping[str, str]] = None): ...
 
     @abstractmethod
     async def close(self) -> None: ...

hyperbrowser/transport/sync.py

@@ -1,6 +1,6 @@
 import json
 import httpx
-from typing import Optional
+from typing import Mapping, Optional
 
 from hyperbrowser.exceptions import HyperbrowserError
 from hyperbrowser.version import __version__
@@ -10,7 +10,7 @@
 class SyncTransport(SyncTransportStrategy):
     """Synchronous transport implementation using httpx"""
 
-    def __init__(self, api_key: str, headers: Optional[dict] = None):
+    def __init__(self, api_key: str, headers: Optional[Mapping[str, str]] = None):
         merged_headers = {
             "x-api-key": api_key,
             "User-Agent": f"hyperbrowser-python-sdk/{__version__}",

tests/test_custom_headers.py

@@ -1,4 +1,5 @@
 import asyncio
+from types import MappingProxyType
 
 import pytest
 
@@ -83,6 +84,33 @@ async def run() -> None:
     asyncio.run(run())
 
 
+def test_sync_client_constructor_accepts_mapping_headers():
+    source_headers = {"X-Team-Trace": "team-mapping-sync"}
+    mapping_headers = MappingProxyType(source_headers)
+    client = Hyperbrowser(api_key="test-key", headers=mapping_headers)
+    source_headers["X-Team-Trace"] = "mutated"
+    try:
+        assert client.transport.client.headers["X-Team-Trace"] == "team-mapping-sync"
+    finally:
+        client.close()
+
+
+def test_async_client_constructor_accepts_mapping_headers():
+    async def run() -> None:
+        source_headers = {"X-Team-Trace": "team-mapping-async"}
+        mapping_headers = MappingProxyType(source_headers)
+        client = AsyncHyperbrowser(api_key="test-key", headers=mapping_headers)
+        source_headers["X-Team-Trace"] = "mutated"
+        try:
+            assert (
+                client.transport.client.headers["X-Team-Trace"] == "team-mapping-async"
+            )
+        finally:
+            await client.close()
+
+    asyncio.run(run())
+
+
 def test_client_constructor_rejects_mixed_config_and_direct_args():
     with pytest.raises(TypeError, match="Pass either `config`"):
         Hyperbrowser(