Disable nested monty decode for some data (#784)

* Disable nested monty decode for some data * Linting * Test fixes * Linting * Update get_entries test * More linting
materialsproject · Apr 27, 2023 · d5190ff · d5190ff
1 parent 6316599
commit d5190ff
Show file tree

Hide file tree

Showing 3 changed files with 9 additions and 5 deletions.
diff --git a/mp_api/client/mprester.py b/mp_api/client/mprester.py
@@ -199,7 +199,9 @@ def __init__(
                 endpoint=endpoint,
                 include_user_agent=include_user_agent,
                 session=self.session,
-                monty_decode=monty_decode,
+                monty_decode=monty_decode
+                if cls not in [TaskRester, ProvenanceRester]  # type: ignore
+                else False,  # Disable monty decode on nested data which may give errors
                 use_document_model=use_document_model,
                 headers=self.headers,
             )  # type: BaseRester
@@ -594,7 +596,8 @@ def get_entries(
                         entry_dict["composition"][element] *= site_ratio
 
                     for correction in entry_dict["energy_adjustments"]:
-                        correction["n_atoms"] *= site_ratio
+                        if "n_atoms" in correction:
+                            correction["n_atoms"] *= site_ratio
 
                 entry = (
                     ComputedStructureEntry.from_dict(entry_dict)

diff --git a/tests/test_mprester.py b/tests/test_mprester.py
@@ -160,9 +160,9 @@ def test_get_entries(self, mpr):
 
         s = prim.structure
         assert pytest.approx(s.lattice.a) == s.lattice.b
-        assert pytest.approx(s.lattice.a) != s.lattice.c
+        assert pytest.approx(s.lattice.a) == s.lattice.c
         assert pytest.approx(s.lattice.alpha) == s.lattice.beta
-        assert pytest.approx(s.lattice.alpha) != s.lattice.gamma
+        assert pytest.approx(s.lattice.alpha) == s.lattice.gamma
 
         # Additional criteria
         entry = mpr.get_entries(

diff --git a/tests/test_thermo.py b/tests/test_thermo.py
@@ -121,5 +121,6 @@ def test_get_phase_diagram_from_chemsys():
     # Test that a phase diagram is returned
 
     assert isinstance(
-        ThermoRester().get_phase_diagram_from_chemsys("Hf-Pm"), PhaseDiagram
+        ThermoRester().get_phase_diagram_from_chemsys("Hf-Pm", thermo_type="GGA_GGA+U"),
+        PhaseDiagram,
     )