merge conflicts

Author: esoteric-ephemera

.gitignore

@@ -120,4 +120,4 @@ ENV/
 _autosummary
 
 uv.lock
-JANAF_O2_data.json
+JANAF_*_data.json

LICENSE

@@ -1,4 +1,4 @@
-maggma Copyright (c) 2017, The Regents of the University of
+Copyright (c) 2017, The Regents of the University of
 California, through Lawrence Berkeley National Laboratory (subject
 to receipt of any required approvals from the U.S. Dept. of Energy).
 All rights reserved.

mp_api/client/core/client.py

@@ -7,7 +7,6 @@
 
 import inspect
 import itertools
-import json
 import os
 import platform
 import sys
@@ -23,15 +22,14 @@
     TYPE_CHECKING,
     ForwardRef,
     Generic,
+    Optional,
     TypeVar,
     get_args,
 )
 from urllib.parse import quote, urljoin
 
 import requests
-from bson import json_util
 from emmet.core.utils import jsanitize
-from monty.json import MontyDecoder
 from pydantic import BaseModel, create_model
 from requests.adapters import HTTPAdapter
 from requests.exceptions import RequestException
@@ -40,7 +38,7 @@
 from urllib3.util.retry import Retry
 
 from mp_api.client.core.settings import MAPIClientSettings
-from mp_api.client.core.utils import api_sanitize, validate_ids
+from mp_api.client.core.utils import load_json, validate_ids
 
 try:
     import boto3
@@ -57,6 +55,8 @@
 if TYPE_CHECKING:
     from typing import Any, Callable
 
+    from pydantic.fields import FieldInfo
+
 try:
     __version__ = version("mp_api")
 except PackageNotFoundError:  # pragma: no cover
@@ -150,12 +150,6 @@ def __init__(
         else:
             self._s3_client = None
 
-        self.document_model = (
-            api_sanitize(self.document_model)  # type: ignore
-            if self.document_model is not None
-            else None  # type: ignore
-        )
-
     @property
     def session(self) -> requests.Session:
         if not self._session:
@@ -270,11 +264,7 @@ def _post_resource(
             response = self.session.post(url, json=payload, verify=True, params=params)
 
             if response.status_code == 200:
-                if self.monty_decode:
-                    data = json.loads(response.text, cls=MontyDecoder)
-                else:
-                    data = json.loads(response.text)
-
+                data = load_json(response.text, deser=self.monty_decode)
                 if self.document_model and use_document_model:
                     if isinstance(data["data"], dict):
                         data["data"] = self.document_model.model_validate(data["data"])  # type: ignore
@@ -287,7 +277,7 @@ def _post_resource(
 
             else:
                 try:
-                    data = json.loads(response.text)["detail"]
+                    data = load_json(response.text)["detail"]
                 except (JSONDecodeError, KeyError):
                     data = f"Response {response.text}"
                 if isinstance(data, str):
@@ -342,11 +332,7 @@ def _patch_resource(
             response = self.session.patch(url, json=payload, verify=True, params=params)
 
             if response.status_code == 200:
-                if self.monty_decode:
-                    data = json.loads(response.text, cls=MontyDecoder)
-                else:
-                    data = json.loads(response.text)
-
+                data = load_json(response.text, deser=self.monty_decode)
                 if self.document_model and use_document_model:
                     if isinstance(data["data"], dict):
                         data["data"] = self.document_model.model_validate(data["data"])  # type: ignore
@@ -359,7 +345,7 @@ def _patch_resource(
 
             else:
                 try:
-                    data = json.loads(response.text)["detail"]
+                    data = load_json(response.text)["detail"]
                 except (JSONDecodeError, KeyError):
                     data = f"Response {response.text}"
                 if isinstance(data, str):
@@ -384,18 +370,24 @@ def _query_open_data(
         self,
         bucket: str,
         key: str,
-        decoder: Callable,
+        decoder: Callable | None = None,
     ) -> tuple[list[dict] | list[bytes], int]:
         """Query and deserialize Materials Project AWS open data s3 buckets.
 
         Args:
             bucket (str): Materials project bucket name
             key (str): Key for file including all prefixes
-            decoder(Callable): Callable used to deserialize data
+            decoder(Callable or None): Callable used to deserialize data.
+                Defaults to mp_api.core.utils.load_json
 
         Returns:
             dict: MontyDecoded data
         """
+        if not decoder:
+
+            def decoder(x):
+                return load_json(x, deser=self.monty_decode)
+
         file = open(
             f"s3://{bucket}/{key}",
             encoding="utf-8",
@@ -527,16 +519,11 @@ def _query_resource(
                         "Ignoring `fields` argument: All fields are always included when no query is provided."
                     )
 
-                decoder = (
-                    MontyDecoder().decode if self.monty_decode else json_util.loads
-                )
-
                 # Multithreaded function inputs
                 s3_params_list = {
                     key: {
                         "bucket": bucket,
                         "key": key,
-                        "decoder": decoder,
                     }
                     for key in keys
                 }
@@ -1013,11 +1000,7 @@ def _submit_request_and_process(
             )
 
         if response.status_code == 200:
-            if self.monty_decode:
-                data = json.loads(response.text, cls=MontyDecoder)
-            else:
-                data = json.loads(response.text)
-
+            data = load_json(response.text, deser=self.monty_decode)
             # other sub-urls may use different document models
             # the client does not handle this in a particularly smart way currently
             if self.document_model and use_document_model:
@@ -1029,7 +1012,7 @@ def _submit_request_and_process(
 
         else:
             try:
-                data = json.loads(response.text)["detail"]
+                data = load_json(response.text)["detail"]
             except (JSONDecodeError, KeyError):
                 data = f"Response {response.text}"
             if isinstance(data, str):
@@ -1057,10 +1040,8 @@ def _convert_to_model(self, data: list[dict]):
             (list[MPDataDoc]): List of MPDataDoc objects
 
         """
-        raw_doc_list = [self.document_model.model_validate(d) for d in data]  # type: ignore
-
-        if len(raw_doc_list) > 0:
-            data_model, set_fields, _ = self._generate_returned_model(raw_doc_list[0])
+        if len(data) > 0:
+            data_model, set_fields, _ = self._generate_returned_model(data[0])
 
             data = [
                 data_model(
@@ -1070,44 +1051,56 @@ def _convert_to_model(self, data: list[dict]):
                         if field in set_fields
                     }
                 )
-                for raw_doc in raw_doc_list
+                for raw_doc in data
             ]
 
         return data
 
-    def _generate_returned_model(self, doc):
+    def _generate_returned_model(
+        self, doc: dict[str, Any]
+    ) -> tuple[BaseModel, list[str], list[str]]:
         model_fields = self.document_model.model_fields
-
-        set_fields = doc.model_fields_set
+        set_fields = [k for k in doc if k in model_fields]
         unset_fields = [field for field in model_fields if field not in set_fields]
 
         # Update with locals() from external module if needed
-        other_vars = {}
         if any(
+            isinstance(field_meta.annotation, ForwardRef)
+            for field_meta in model_fields.values()
+        ) or any(
             isinstance(typ, ForwardRef)
             for field_meta in model_fields.values()
             for typ in get_args(field_meta.annotation)
         ):
-            other_vars = vars(import_module(self.document_model.__module__))
-
-        include_fields = {
-            name: (
-                model_fields[name].annotation,
-                model_fields[name],
+            vars(import_module(self.document_model.__module__))
+
+        include_fields: dict[str, tuple[type, FieldInfo]] = {}
+        for name in set_fields:
+            field_copy = model_fields[name]._copy()
+            field_copy.default = None
+            include_fields[name] = (
+                Optional[model_fields[name].annotation],
+                field_copy,
             )
-            for name in set_fields
-        }
 
         data_model = create_model(  # type: ignore
             "MPDataDoc",
             **include_fields,
             # TODO fields_not_requested is not the same as unset_fields
             # i.e. field could be requested but not available in the raw doc
             fields_not_requested=(list[str], unset_fields),
-            __base__=self.document_model,
+            __doc__=".".join(
+                [
+                    getattr(self.document_model, k, "")
+                    for k in ("__module__", "__name__")
+                ]
+            ),
+            __module__=self.document_model.__module__,
         )
-        if other_vars:
-            data_model.model_rebuild(_types_namespace=other_vars)
+        # if other_vars:
+        #     data_model.model_rebuild(_types_namespace=other_vars)
+
+        orig_rester_name = self.document_model.__name__
 
         def new_repr(self) -> str:
             extra = ",\n".join(
@@ -1116,7 +1109,7 @@ def new_repr(self) -> str:
                 if n == "fields_not_requested" or n in set_fields
             )
 
-            s = f"\033[4m\033[1m{self.__class__.__name__}<{self.__class__.__base__.__name__}>\033[0;0m\033[0;0m(\n{extra}\n)"  # noqa: E501
+            s = f"\033[4m\033[1m{self.__class__.__name__}<{orig_rester_name}>\033[0;0m\033[0;0m(\n{extra}\n)"  # noqa: E501
             return s
 
         def new_str(self) -> str:
@@ -1230,8 +1223,14 @@ def get_data_by_id(
             stacklevel=2,
         )
 
-        if self.primary_key in ["material_id", "task_id"]:
-            validate_ids([document_id])
+        if self.primary_key in [
+            "material_id",
+            "task_id",
+            "battery_id",
+            "spectrum_id",
+            "thermo_id",
+        ]:
+            document_id = validate_ids([document_id])[0]
 
         if isinstance(fields, str):  # pragma: no cover
             fields = (fields,)  # type: ignore