move logic for setting consumer/api header stuff to mprester

Author: esoteric-ephemera

mp_api/client/_server_utils.py

@@ -1,4 +1,5 @@
-"""Define utilities needed by the MP web server."""
+"""Define flask-dependent utilities for the web server."""
+
 from __future__ import annotations
 
 try:
@@ -10,7 +11,6 @@
     _has_request_context = None
     request = None
 
-from mp_api.client import MPRester
 from mp_api.client.core.utils import validate_api_key
 
 def has_request_context() -> bool:
@@ -34,8 +34,8 @@ def get_request_headers() -> dict[str,Any]:
     """
     return request.headers if has_request_context() else {}
 
-def is_localhost() -> bool:
-    """Determine if current env is local or production.
+def is_dev_env() -> bool:
+    """Determine if current env is local/developmental or production.
 
     Returns:
         bool: True if the environment is locally hosted.
@@ -83,39 +83,24 @@ def is_logged_in_user(consumer: dict[str, str] | None = None) -> bool:
     return bool(not c.get("X-Anonymous-Consumer") and c.get("X-Consumer-Id"))
 
 
-def get_user_api_key(consumer: dict[str, str] | None = None) -> str | None:
+def get_user_api_key(
+    api_key : str | None = None,
+    consumer: dict[str, str] | None = None
+) -> str | None:
     """Get the api key that belongs to the current user.
 
     If running on localhost, api key is obtained from
     the environment variable MP_API_KEY.
 
     Args:
+        api_key (str or None) : User API key
         consumer (dict of str to str, or None): Headers associated with the consumer
 
     Returns:
         str, the API key, or None if no API key could be identified.
     """
-    c = consumer or get_consumer()
-
-    if is_localhost():
-        return validate_api_key()
-    elif is_logged_in_user(c):
+    if is_dev_env():
+        return validate_api_key(api_key=api_key)
+    elif is_logged_in_user(c := consumer or get_consumer()):
         return c.get("X-Consumer-Custom-Id")
-    return None
-
-
-def get_rester(**kwargs) -> MPRester:
-    """Create MPRester with headers set for localhost and production compatibility.
-
-    Args:
-        **kwargs : kwargs to pass to MPRester
-
-    Returns:
-        MPRester
-    """
-    if is_localhost():
-        dev_api_key = get_user_api_key()
-        SESSION.headers["x-api-key"] = dev_api_key or ""
-        return MPRester(api_key=dev_api_key, session=SESSION, **kwargs)
-
-    return MPRester(headers=get_consumer(), session=SESSION, **kwargs)
+    return None

mp_api/client/mprester.py

@@ -21,6 +21,7 @@
 from pymatgen.symmetry.analyzer import SpacegroupAnalyzer
 from requests import Session, get
 
+from mp_api.client._server_utils import get_consumer, get_user_api_key, is_dev_env
 from mp_api.client.core import BaseRester
 from mp_api.client.core._oxygen_evolution import OxygenEvolution
 from mp_api.client.core.exceptions import (
@@ -32,7 +33,6 @@
 from mp_api.client.core.utils import (
     LazyImport,
     load_json,
-    validate_api_key,
     validate_endpoint,
     validate_ids,
 )
@@ -141,16 +141,18 @@ def __init__(
             force_renew: Option to overwrite existing local dataset
             **kwargs: access to legacy kwargs that may be in the process of being deprecated
         """
-        self.api_key = validate_api_key(api_key)
+        self.api_key = get_user_api_key(api_key=api_key)
 
         self.endpoint = validate_endpoint(endpoint)
 
-        self.headers = headers or {}
+        self.headers = headers or get_consumer()
         self.session = session or BaseRester._create_session(
             api_key=self.api_key,
             include_user_agent=include_user_agent,
             headers=self.headers,
         )
+        if is_dev_env():
+            self.session.headers["x-api-key"] = self.api_key
         self._include_user_agent = include_user_agent
         self.use_document_model = use_document_model
         self.mute_progress_bars = mute_progress_bars

tests/client/core/test_utils.py

@@ -142,15 +142,15 @@ def test_api_key_validation(monkeypatch: pytest.MonkeyPatch):
     monkeypatch.setattr(pymatgen.core, "SETTINGS", non_api_key_settings)
 
     with pytest.raises(MPRestError, match="32 characters"):
-        validate_api_key("invalid_key")
+        validate_api_key(api_key="invalid_key")
 
     with pytest.warns(MPRestWarning, match="No API key found"):
         validate_api_key()
 
     junk_api_key = "a" * 32
     monkeypatch.setenv("MP_API_KEY", junk_api_key)
     assert validate_api_key() == junk_api_key
-    assert validate_api_key(junk_api_key) == junk_api_key
+    assert validate_api_key(api_key=junk_api_key) == junk_api_key
 
     other_junk_api_key = "b" * 32
     monkeypatch.setattr(