fix re pattern

materialsproject · Dec 24, 2024 · d714101 · d714101
1 parent 3474e79
commit d714101
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 6 deletions.
diff --git a/src/pymatgen/analysis/pourbaix_diagram.py b/src/pymatgen/analysis/pourbaix_diagram.py
@@ -411,14 +411,14 @@ def ion_or_solid_comp_object(formula: str) -> Composition | Ion:
         Composition/Ion object
     """
     # Formula for ion
-    try:
+    if formula.endswith("(aq)") or re.search(r"\[.*\]$", formula) or "-" in formula or "+" in formula:
         return Ion.from_formula(formula)
 
     # Formula for solid
-    except ValueError:
-        if formula.endswith("(s)"):
-            return Composition(formula[:-3])
-        return Composition(formula)
+    if formula.endswith("(s)"):
+        return Composition(formula[:-3])
+
+    return Composition(formula)
 
 
 # TODO: the solids filter breaks some of the functionality of the

diff --git a/tests/analysis/test_pourbaix_diagram.py b/tests/analysis/test_pourbaix_diagram.py
@@ -335,7 +335,6 @@ def test_ion(self):
         assert ion_or_solid_comp_object("Ca[++]").charge == 2
         assert ion_or_solid_comp_object("Ca2+").charge == 1
         assert ion_or_solid_comp_object("C2O4-2").charge == -2
-        assert ion_or_solid_comp_object("CO2").charge == 0
 
         # Test anions
         assert ion_or_solid_comp_object("Cl-").charge == -1