Add document models, type hints

src/mp_api/piezo/client.py

@@ -2,13 +2,15 @@
 from collections import defaultdict
 
 from mp_api.core.client import BaseRester, MPRestError
+from mp_api.piezo.models import PiezoDoc
 
 import warnings
 
 
 class PiezoRester(BaseRester):
 
     suffix = "piezoelectric"
+    document_model = PiezoDoc
 
     def get_piezo_from_material_id(self, material_id: str):
         """

src/mp_api/piezo/client.pyi

@@ -0,0 +1,14 @@
+from typing import List, Optional
+from mp_api.piezo.models import PiezoDoc
+
+
+class PiezoRester:
+
+    def get_document_by_id(
+            self,
+            document_id: str,
+            fields: Optional[List[str]] = None,
+            monty_decode: bool = True,
+            version: Optional[str] = None,
+    ) -> PiezoDoc:
+        ...

src/mp_api/search/client.py

@@ -3,5 +3,6 @@
 
 
 class SearchRester(BaseRester):
+
     suffix = "search"
     document_model = SearchDoc

src/mp_api/search/client.pyi

@@ -0,0 +1,14 @@
+from typing import List, Optional
+from mp_api.search.models import SearchDoc
+
+
+class SearchRester:
+
+    def get_document_by_id(
+            self,
+            document_id: str,
+            fields: Optional[List[str]] = None,
+            monty_decode: bool = True,
+            version: Optional[str] = None,
+    ) -> SearchDoc:
+        ...

src/mp_api/substrates/client.py

@@ -3,11 +3,13 @@
 import warnings
 
 from mp_api.core.client import BaseRester
+from mp_api.substrates.models import SubstratesDoc
 
 
 class SubstratesRester(BaseRester):
 
     suffix = "substrates"
+    document_model = SubstratesDoc
 
     def search_substrates_docs(
         self,

src/mp_api/substrates/client.pyi

@@ -0,0 +1,14 @@
+from typing import List, Optional
+from mp_api.substrates.models import SubstratesDoc
+
+
+class SubstratesRester:
+
+    def get_document_by_id(
+            self,
+            document_id: str,
+            fields: Optional[List[str]] = None,
+            monty_decode: bool = True,
+            version: Optional[str] = None,
+    ) -> SubstratesDoc:
+        ...

src/mp_api/surface_properties/client.py

@@ -3,11 +3,13 @@
 import warnings
 
 from mp_api.core.client import BaseRester, MPRestError
+from mp_api.surface_properties.models import SurfacePropDoc
 
 
 class SurfacePropertiesRester(BaseRester):
 
     suffix = "surface_properties"
+    document_model = SurfacePropDoc
 
     def get_surface_properties_from_material_id(self, material_id: str):
         """

src/mp_api/surface_properties/client.pyi

@@ -0,0 +1,14 @@
+from typing import List, Optional
+from mp_api.surface_properties.models import SurfacePropDoc
+
+
+class SurfacePropertiesRester:
+
+    def get_document_by_id(
+            self,
+            document_id: str,
+            fields: Optional[List[str]] = None,
+            monty_decode: bool = True,
+            version: Optional[str] = None,
+    ) -> SurfacePropDoc:
+        ...

src/mp_api/synthesis/client.pyi

@@ -3,5 +3,6 @@ from mp_api.synthesis.models import SynthesisDoc
 
 
 class SynthesisRester:
+
     def query_text(self, keywords: List[str]) -> SynthesisDoc:
         ...

src/mp_api/tasks/client.pyi

@@ -3,7 +3,12 @@ from mp_api.tasks.models import TaskDoc
 
 
 class TaskRester:
-    def query_by_task_id(
-        self, task_id, fields: Optional[List[str]] = None, monty_decode: bool = True
+
+    def get_document_by_id(
+            self,
+            document_id: str,
+            fields: Optional[List[str]] = None,
+            monty_decode: bool = True,
+            version: Optional[str] = None,
     ) -> TaskDoc:
         ...

src/mp_api/thermo/client.py

@@ -11,6 +11,7 @@ class ThermoRester(BaseRester):
 
     suffix = "thermo"
     document_model = ThermoDoc
+    supports_versions = True
 
     def search_thermo_docs(
         self,

src/mp_api/thermo/client.pyi

@@ -1,9 +1,13 @@
 from typing import List, Optional
 from mp_api.thermo.models import ThermoDoc
 
-
 class ThermoRester:
-    def query_by_task_id(
-        self, task_id, fields: Optional[List[str]] = None, monty_decode: bool = True
+
+    def get_document_by_id(
+            self,
+            document_id: str,
+            fields: Optional[List[str]] = None,
+            monty_decode: bool = True,
+            version: Optional[str] = None,
     ) -> ThermoDoc:
         ...

src/mp_api/xas/client.py

@@ -1,12 +1,13 @@
 from typing import List, Optional
 from pymatgen.core.periodic_table import Element
 from mp_api.core.client import BaseRester, MPRestError
-from mp_api.xas.models import Edge, XASType
+from mp_api.xas.models import Edge, XASType, XASDoc
 
 
 class XASRester(BaseRester):
 
     suffix = "xas"
+    document_model = XASDoc
 
     def get_available_elements(
         self,

src/mp_api/xas/client.pyi

@@ -0,0 +1,14 @@
+from typing import List, Optional
+from mp_api.xas.models import XASDoc
+
+
+class XASRester:
+
+    def get_document_by_id(
+            self,
+            document_id: str,
+            fields: Optional[List[str]] = None,
+            monty_decode: bool = True,
+            version: Optional[str] = None,
+    ) -> XASDoc:
+        ...