Skip to content

Commit

Permalink
Deprecate get_data_by_id (#882)
Browse files Browse the repository at this point in the history
* Remove get_data_by_id

* Add search method for DOI rester

* Big endpoint updates

* Remove charge density rester

* Fix materials structure method

* Remove charge density references

* More charge density fixes

* Add type ignore

* Fix task_id method

* Fix references method

* Fix phonon methods

* Fix phonon and task methods

* Fix remaining type issues in MPRester

* Remove get_bv refs in es methods

* Fix user settings method

* Fix get_by ref in molecules

* Remove remaining references to get_data_by_id

* Deprecate get_data_by_id

* Linting

* Remove charge density rester tests

* Remove generic get method tests

* Remove warning expectation

* Linting

* Fix method reference in test

* More method typo fixes

* More materials typos

* Linting

* Fix boto3 warning

* Fix generic search tests

* Linting

* Update task_ids arg

* Fix remaining search methods

* Fix thermo and tasks

* Linting
  • Loading branch information
munrojm committed Feb 14, 2024
1 parent 8a84e76 commit 96cf642
Show file tree
Hide file tree
Showing 39 changed files with 566 additions and 739 deletions.
110 changes: 34 additions & 76 deletions mp_api/client/core/client.py
Expand Up @@ -128,7 +128,9 @@ def __init__(
self._s3_resource = None

self.document_model = (
api_sanitize(self.document_model) if self.document_model is not None else None # type: ignore
api_sanitize(self.document_model)
if self.document_model is not None
else None # type: ignore
)

@property
Expand Down Expand Up @@ -237,7 +239,9 @@ def _post_resource(
if isinstance(data["data"], dict):
data["data"] = self.document_model.parse_obj(data["data"]) # type: ignore
elif isinstance(data["data"], list):
data["data"] = [self.document_model.parse_obj(d) for d in data["data"]] # type: ignore
data["data"] = [
self.document_model.parse_obj(d) for d in data["data"]
] # type: ignore

return data

Expand Down Expand Up @@ -307,7 +311,9 @@ def _patch_resource(
if isinstance(data["data"], dict):
data["data"] = self.document_model.parse_obj(data["data"]) # type: ignore
elif isinstance(data["data"], list):
data["data"] = [self.document_model.parse_obj(d) for d in data["data"]] # type: ignore
data["data"] = [
self.document_model.parse_obj(d) for d in data["data"]
] # type: ignore

return data

Expand Down Expand Up @@ -967,79 +973,6 @@ def _query_resource_data(
num_chunks=1,
).get("data")

def get_data_by_id(
self,
document_id: str,
fields: list[str] | None = None,
) -> T:
"""Query the endpoint for a single document.
Arguments:
document_id: the unique key for this kind of document, typically a task_id
fields: list of fields to return, by default will return all fields
Returns:
A single document.
"""
if document_id is None:
raise ValueError(
"Please supply a specific ID. You can use the query method to find "
"ids of interest."
)

if self.primary_key in ["material_id", "task_id"]:
validate_ids([document_id])

if fields is None:
criteria = {"_all_fields": True, "_limit": 1} # type: dict
else:
criteria = {"_limit": 1}

if isinstance(fields, str): # pragma: no cover
fields = (fields,)

results = [] # type: List

try:
results = self._query_resource_data(criteria=criteria, fields=fields, suburl=document_id) # type: ignore
except MPRestError:
if self.primary_key == "material_id":
# see if the material_id has changed, perhaps a task_id was supplied
# this should likely be re-thought
from mp_api.client.routes.materials.materials import MaterialsRester

with MaterialsRester(
api_key=self.api_key,
endpoint=self.base_endpoint,
use_document_model=False,
monty_decode=False,
session=self.session,
headers=self.headers,
) as mpr:
docs = mpr.search(task_ids=[document_id], fields=["material_id"])

if len(docs) > 0:
new_document_id = docs[0].get("material_id", None)

if new_document_id is not None:
warnings.warn(
f"Document primary key has changed from {document_id} to {new_document_id}, "
f"returning data for {new_document_id} in {self.suffix} route. "
)

results = self._query_resource_data(
criteria=criteria, fields=fields, suburl=new_document_id # type: ignore
)

if not results:
raise MPRestError(f"No result for record {document_id}.")
elif len(results) > 1: # pragma: no cover
raise ValueError(
f"Multiple records for {document_id}, this shouldn't happen. Please report as a bug."
)
else:
return results[0]

def _search(
self,
num_chunks: int | None = None,
Expand Down Expand Up @@ -1078,6 +1011,31 @@ def _search(
num_chunks=num_chunks,
)

def get_data_by_id(
self,
document_id: str,
fields: list[str] | None = None,
) -> T | dict:
warnings.warn(
"get_data_by_id is deprecated and will be removed soon. Please use the search method instead.",
DeprecationWarning,
stacklevel=2,
)

if self.primary_key in ["material_id", "task_id"]:
validate_ids([document_id])

if isinstance(fields, str): # pragma: no cover
fields = (fields,) # type: ignore

return self._search( # type: ignore
**{self.primary_key + "s": document_id},
num_chunks=1,
chunk_size=1,
all_fields=fields is None,
fields=fields,
)

def _get_all_documents(
self,
query_params,
Expand Down

0 comments on commit 96cf642

Please sign in to comment.