Allow solid_compat to be None in get_pourbaix_entries (#1061)

Author: esoteric-ephemera

mp_api/client/mprester.py

@@ -39,6 +39,7 @@
 
     from emmet.core.tasks import CoreTaskDoc
     from pymatgen.analysis.phase_diagram import PDEntry
+    from pymatgen.entries.compatibility import Compatibility
     from pymatgen.entries.computed_entries import ComputedEntry
 
 
@@ -611,7 +612,11 @@ def get_entries(
     def get_pourbaix_entries(
         self,
         chemsys: str | list[str] | list[ComputedEntry | ComputedStructureEntry],
-        solid_compat="MaterialsProject2020Compatibility",
+        solid_compat: Literal[
+            "MaterialsProjectCompatibility", "MaterialsProject2020Compatibility"
+        ]
+        | Compatibility
+        | None = "MaterialsProject2020Compatibility",
         use_gibbs: Literal[300] | None = None,
     ):
         """A helper function to get all entries necessary to generate
@@ -627,10 +632,12 @@ def get_pourbaix_entries(
                 for adding extra calculation data to the Pourbaix Diagram.
                 If this is set, the chemsys will be inferred from the entries.
             solid_compat: Compatibility scheme used to pre-process solid DFT energies prior
-                to applying aqueous energy adjustments. May be passed as a class (e.g.
-                MaterialsProject2020Compatibility) or an instance
-                (e.g., MaterialsProject2020Compatibility()). If None, solid DFT energies
-                are used as-is. Default: MaterialsProject2020Compatibility
+                to applying aqueous energy adjustments.
+                May be passed as a string (either "MaterialsProjectCompatibility"
+                or "MaterialsProject2020Compatibility"), or as a class instance
+                (e.g., MaterialsProject2020Compatibility()).
+                If None, solid DFT energies are used as-is.
+                Default: MaterialsProject2020Compatibility
             use_gibbs: Set to 300 (for 300 Kelvin) to use a machine learning model to
                 estimate solid free energy from DFT energy (see GibbsComputedStructureEntry).
                 This can slightly improve the accuracy of the Pourbaix diagram in some
@@ -673,12 +680,12 @@ def get_pourbaix_entries(
             solid_compat = MaterialsProjectCompatibility()
         elif solid_compat == "MaterialsProject2020Compatibility":
             solid_compat = MaterialsProject2020Compatibility()
-        elif isinstance(solid_compat, Compatibility):
+        elif isinstance(solid_compat, Compatibility) or solid_compat is None:
             pass
         else:
             raise ValueError(
                 "Solid compatibility can only be 'MaterialsProjectCompatibility', "
-                "'MaterialsProject2020Compatibility', or an instance of a Compatibility class"
+                "'MaterialsProject2020Compatibility', None, or an instance of a Compatibility class"
             )
 
         pbx_entries = []

tests/client/test_mprester.py

@@ -270,7 +270,7 @@ def test_get_pourbaix_entries(self, mpr):
 
         # test solid_compat kwarg
         with pytest.raises(ValueError, match="Solid compatibility can only be"):
-            mpr.get_pourbaix_entries("Ti-O", solid_compat=None)
+            mpr.get_pourbaix_entries("Ti-O", solid_compat="None")
 
         # test removal of extra elements from reference solids
         # Li-Zn-S has Na in reference solids