move unset fields to separate attr

Author: esoteric-ephemera

mp_api/client/core/client.py

@@ -1205,12 +1205,14 @@ def _submit_request_and_process(
             # 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:
-                requested_fields = (
-                    params.get("_fields", "").split(",")
-                    if params.get("_fields")
-                    else None
+                data["data"] = self._convert_to_model(
+                    data["data"],
+                    requested_fields=(
+                        params["_fields"].split(",")
+                        if isinstance(params.get("_fields"), str)
+                        else None
+                    ),
                 )
-                data["data"] = self._convert_to_model(data["data"], requested_fields)
 
             meta_total_doc_num = data.get("meta", {}).get("total_doc", 1)
 
@@ -1260,15 +1262,14 @@ def _convert_to_model(
                 # Return empty list if no data in iterator
                 return []
             data_model, set_fields, _ = self._generate_returned_model(
-                first_doc, requested_fields
+                first_doc, requested_fields=requested_fields
             )
 
             return [
                 data_model(
                     **{
-                        field: value
-                        for field, value in dict(raw_doc).items()
-                        if field in set_fields
+                        field: raw_doc[field]
+                        for field in set_fields.intersection(raw_doc)
                     }
                 )
                 for raw_doc in (data if is_list else chain([first_doc], data))
@@ -1280,24 +1281,23 @@ def _generate_returned_model(
         self,
         doc: dict[str, Any],
         requested_fields: list[str] | None = None,
-    ) -> tuple[type[BaseModel], list[str], list[str]]:
+    ) -> tuple[type[BaseModel], set[str], set[str]]:
         """Dynamically generates an MPDataDoc Pydantic model from API response content.
 
         Args:
             doc (dict): A single document returned from the API
-            requested_fields (list[str] or None): Optional list of fields to be returned
+            requested_fields (list of str, or None): Optional list of fields to be returned
 
         Returns:
-            (tuple): Tuple containing (data_model, set_fields, fields_not_requested)
+            BaseModel: the pydantic model representing the data
+            set of str: fields set in the document model
+            set of str: set_fields, fields_not_requested)
         """
         model_fields = self.document_model.model_fields
-        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]
-        fields_not_requested = (
-            [field for field in unset_fields if field not in requested_fields]
-            if requested_fields
-            else unset_fields
-        )
+        set_fields = set(doc).intersection(model_fields)
+        unset_fields = set(model_fields).difference(set_fields)
+        user_requested_fields: list[str] = requested_fields or []
+        fields_not_requested = unset_fields.difference(user_requested_fields)
 
         # Update with locals() from external module if needed
         if any(
@@ -1324,7 +1324,11 @@ def _generate_returned_model(
         data_model = create_model(  # type: ignore
             "MPDataDoc",
             **include_fields,
-            fields_not_requested=(list[str], fields_not_requested),
+            fields_not_requested=(list[str], list(fields_not_requested)),
+            unavailable_fields=(
+                list[str],
+                list(unset_fields.intersection(user_requested_fields)),
+            ),
             __base__=_DictLikeAccess,
             __doc__=".".join(
                 [
@@ -1354,13 +1358,19 @@ def new_str(self) -> str:
                 if 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\n\033[1mFields not requested:\033[0;0m\n{fields_not_requested}"  # noqa: E501
-            return s
+            return (
+                f"\033[4m\033[1m{self.__class__.__name__}"
+                f"<{self.__class__.__base__.__name__}>\033[0;0m\033[0;0m"
+                f"\n{extra}\n\n"
+                f"\033[1mFields not requested:\033[0;0m\n{fields_not_requested}"
+            )
 
         def new_getattr(self, attr) -> str:
+            if attr in self.unavailable_fields:
+                raise AttributeError(f"`{attr}` is unavailable in the returned data.")
             if attr in self.fields_not_requested:
                 raise AttributeError(
-                    f"'{attr}' data is available but has not been requested in 'fields'."
+                    f"`{attr}` data is available but has not been requested in `fields`."
                     " A full list of unrequested fields can be found in `fields_not_requested`."
                 )
             else:

tests/client/test_core_client.py

@@ -47,15 +47,27 @@ def test_available_fields(rester, mpr):
     assert rester.available_fields == []
 
 
-def test_fields_not_requested_excludes_requested_fields(rester: BaseRester):
-    from emmet.core.tasks import TaskDoc
+def test_fields_not_requested_excludes_requested_fields(mpr):
+    task_rester = mpr.materials.tasks
+    doc = {"task_id": "fakeid-1234", "state": "successful"}
+    requested_fields = list(task_rester.document_model.model_fields)
 
-    rester.document_model = TaskDoc
-    doc = {"task_id": "fakeid-1234"}
-    requested_fields = list(TaskDoc.model_fields.keys())
-
-    _, _, fields_not_requested = rester._generate_returned_model(
+    doc_model, _, fields_not_requested = task_rester._generate_returned_model(
         doc, requested_fields=requested_fields
     )
+    deser_doc = doc_model(**doc)
+
+    assert "dir_name" not in fields_not_requested
+    assert "tags" in deser_doc.unavailable_fields
+    with pytest.raises(AttributeError, match="is unavailable in the returned data"):
+        deser_doc.tags
 
-    assert fields_not_requested == []
+    requested_fields.remove("tags")
+    doc_model, _, _ = task_rester._generate_returned_model(
+        doc, requested_fields=requested_fields
+    )
+    deser_doc = doc_model(**doc)
+    with pytest.raises(
+        AttributeError, match="data is available but has not been requested in"
+    ):
+        deser_doc.tags