From 7f69553298c215169b22101c581567e8f4641c8e Mon Sep 17 00:00:00 2001 From: "Haoyu (Daniel)" Date: Tue, 24 Dec 2024 16:26:00 +0800 Subject: [PATCH] add test and try-except fix --- src/pymatgen/analysis/pourbaix_diagram.py | 17 ++++----- tests/analysis/test_pourbaix_diagram.py | 42 ++++++++++++++++++++++- 2 files changed, 48 insertions(+), 11 deletions(-) diff --git a/src/pymatgen/analysis/pourbaix_diagram.py b/src/pymatgen/analysis/pourbaix_diagram.py index 99193ce3d15..9d7d968caf5 100644 --- a/src/pymatgen/analysis/pourbaix_diagram.py +++ b/src/pymatgen/analysis/pourbaix_diagram.py @@ -411,17 +411,14 @@ def ion_or_solid_comp_object(formula: str) -> Composition | Ion: Composition/Ion object """ # Formula for ion - if re.match(r"\[([^\[\]]+)\]|\(aq\)", formula): - comp_obj = Ion.from_formula(formula) + try: + return Ion.from_formula(formula) # Formula for solid - elif re.search(r"\(s\)", formula): - comp_obj = Composition(formula[:-3]) - - else: - comp_obj = Composition(formula) - - return comp_obj + except ValueError: + if formula.endswith("(s)"): + return Composition(formula[:-3]) + return Composition(formula) # TODO: the solids filter breaks some of the functionality of the @@ -973,7 +970,7 @@ def __init__(self, pourbaix_diagram: PourbaixDiagram) -> None: """ self._pbx = pourbaix_diagram - def show(self, *args, **kwargs): + def show(self, *args, **kwargs) -> None: """Show the Pourbaix plot. Args: diff --git a/tests/analysis/test_pourbaix_diagram.py b/tests/analysis/test_pourbaix_diagram.py index 22066499257..425cbe23a51 100644 --- a/tests/analysis/test_pourbaix_diagram.py +++ b/tests/analysis/test_pourbaix_diagram.py @@ -8,7 +8,14 @@ from monty.serialization import dumpfn, loadfn from pytest import approx -from pymatgen.analysis.pourbaix_diagram import IonEntry, MultiEntry, PourbaixDiagram, PourbaixEntry, PourbaixPlotter +from pymatgen.analysis.pourbaix_diagram import ( + IonEntry, + MultiEntry, + PourbaixDiagram, + PourbaixEntry, + PourbaixPlotter, + ion_or_solid_comp_object, +) from pymatgen.core.composition import Composition from pymatgen.core.ion import Ion from pymatgen.entries.computed_entries import ComputedEntry @@ -315,3 +322,36 @@ def test_plot_entry_stability(self): binary_plotter = PourbaixPlotter(pd_binary) ax = binary_plotter.plot_entry_stability(self.test_data["Ag-Te"][53]) assert isinstance(ax, plt.Axes) + + +class TestIonOrSolidCompObject: + def test_ion(self): + # Test cations + assert ion_or_solid_comp_object("Li+").charge == 1 + assert ion_or_solid_comp_object("Li[+]").charge == 1 + assert ion_or_solid_comp_object("Ca[2+]").charge == 2 + assert ion_or_solid_comp_object("Ca[+2]").charge == 2 + assert ion_or_solid_comp_object("Ca++").charge == 2 + 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 + assert ion_or_solid_comp_object("Cl[-]").charge == -1 + assert ion_or_solid_comp_object("SO4[-2]").charge == -2 + assert ion_or_solid_comp_object("SO4-2").charge == -2 + assert ion_or_solid_comp_object("SO42-").charge == -1 + assert ion_or_solid_comp_object("SO4--").charge == -2 + assert ion_or_solid_comp_object("SO4[--]").charge == -2 + assert ion_or_solid_comp_object("N3-").charge == -1 + + def test_solid(self): + # Test end with "(s)" + assert ion_or_solid_comp_object("Fe2O3(s)") == Composition("Fe2O3") + assert ion_or_solid_comp_object("Fe(s)") == Composition("Fe1") + + # Test end without "(s)" + # TODO: Na2O should be Composition instead of Ion? + assert ion_or_solid_comp_object("Na2O") == Ion.from_formula("Na2O")