From a6b9f8d7b883459b86fadb94712bf92ac8c3dbbc Mon Sep 17 00:00:00 2001
From: Nick Frasser <1693461+nfrasser@users.noreply.github.com>
Date: Wed, 18 Dec 2024 11:15:10 -0500
Subject: [PATCH] fix(spm): correct api client header handling

---
 cryosparc/api.py  | 6 ++++--
 cryosparc/api.pyi | 7 ++++++-
 2 files changed, 10 insertions(+), 3 deletions(-)

diff --git a/cryosparc/api.py b/cryosparc/api.py
index 09a94f9b..dba3587f 100644
--- a/cryosparc/api.py
+++ b/cryosparc/api.py
@@ -75,9 +75,9 @@ def _construct_request(self, _path: str, _schema, *args, **kwargs) -> Tuple[str,
             elif param_in == "query" and param_name in kwargs:
                 # query param must be in kwargs
                 query_params[param_name] = kwargs.pop(param_name)
-            elif param_in == "header" and param_name in kwargs:
+            elif param_in == "header" and (header_name := param_name.replace("-", "_")) in kwargs:
                 # header must be in kwargs
-                headers[param_name] = kwargs.pop(param_name)
+                headers[param_name] = kwargs.pop(header_name)
             elif param_in == "header" and param_name in client_headers:
                 pass  # in default headers, no action required
             elif param_schema["required"]:
@@ -231,6 +231,7 @@ def __init__(
         base_url: Optional[str] = None,
         *,
         auth: Optional[Auth] = None,  # token or email/password
+        headers: Dict[str, str] | None = None,
         timeout: float = 300,
         http_client: Optional[httpx.Client] = None,
     ):
@@ -240,6 +241,7 @@ def __init__(
             if not base_url:
                 raise TypeError(f"Must specify either base_url ({base_url}) or http_client ({http_client})")
             http_client = httpx.Client(base_url=base_url, timeout=timeout)
+        http_client.headers.update(headers)
         super().__init__(http_client)
         self._attrs = set()
         self(auth=auth)  # query the OpenAPI server and populate endpoint functions
diff --git a/cryosparc/api.pyi b/cryosparc/api.pyi
index 26aae670..58e42369 100644
--- a/cryosparc/api.pyi
+++ b/cryosparc/api.pyi
@@ -1190,7 +1190,12 @@ class APIClient:
     developer: DeveloperNamespace
 
     def __init__(
-        self, base_url: Optional[str] = None, *, timeout: float = ..., auth: Union[str, tuple[str, str], None] = None
+        self,
+        base_url: Optional[str] = None,
+        *,
+        auth: Union[str, tuple[str, str], None] = None,
+        headers: Optional[Dict[str, str]] = None,
+        timeout: float = ...,
     ) -> None: ...
     def __call__(self, *, auth: Union[str, tuple[str, str], None] = None) -> Any: ...
     def read_root(self) -> Hello: ...