fix: exclude requested fields from fields_not_requested (#1066)

Author: esoteric-ephemera

mp_api/client/core/client.py

@@ -1205,7 +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:
-                data["data"] = self._convert_to_model(data["data"])
+                data["data"] = self._convert_to_model(
+                    data["data"],
+                    requested_fields=(
+                        params["_fields"].split(",")
+                        if isinstance(params.get("_fields"), str)
+                        else None
+                    ),
+                )
 
             meta_total_doc_num = data.get("meta", {}).get("total_doc", 1)
 
@@ -1234,11 +1241,13 @@ def _submit_request_and_process(
     def _convert_to_model(
         self,
         data: list[dict[str, Any]] | Iterator,
+        requested_fields: list[str] | None = None,
     ) -> list[BaseModel] | list[dict[str, Any]]:
         """Converts dictionary documents to instantiated MPDataDoc objects.
 
         Args:
             data (list[dict] or Iterator): Raw dictionary data objects
+            requested_fields (list[str] or None): Optional list of fields to be returned
 
         Returns:
             (list[MPDataDoc]): List of MPDataDoc objects
@@ -1252,14 +1261,15 @@ def _convert_to_model(
             except StopIteration:
                 # Return empty list if no data in iterator
                 return []
-            data_model, set_fields, _ = self._generate_returned_model(first_doc)
+            data_model, set_fields, _ = self._generate_returned_model(
+                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))
@@ -1268,11 +1278,26 @@ def _convert_to_model(
         return data
 
     def _generate_returned_model(
-        self, doc: dict[str, Any]
-    ) -> tuple[type[BaseModel], list[str], list[str]]:
+        self,
+        doc: dict[str, Any],
+        requested_fields: list[str] | None = None,
+    ) -> 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 of str, or None): Optional list of fields to be returned
+
+        Returns:
+            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]
+        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(
@@ -1299,9 +1324,11 @@ def _generate_returned_model(
         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),
+            fields_not_requested=(list[str], list(fields_not_requested)),
+            unavailable_fields=(
+                list[str],
+                list(unset_fields.intersection(user_requested_fields)),
+            ),
             __base__=_DictLikeAccess,
             __doc__=".".join(
                 [
@@ -1331,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{unset_fields}"  # 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:
@@ -1354,7 +1387,7 @@ def new_dict(self, *args, **kwargs):
         data_model.__getattr__ = new_getattr
         data_model.dict = new_dict
 
-        return data_model, set_fields, unset_fields
+        return data_model, set_fields, fields_not_requested
 
     def _query_resource_data(
         self,

tests/client/test_core_client.py

@@ -45,3 +45,29 @@ def test_count(mpr):
 def test_available_fields(rester, mpr):
     assert len(mpr.materials.available_fields) > 0
     assert rester.available_fields == []
+
+
+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)
+
+    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
+
+    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