Better Composition repr (#3182)

* include cls_name and oxidation states in Composition.__repr__

* delete doctest.testmod() if __name__ == "__main__"

* add tests for Composition.__repr__ and __str__

* merge new tests into test_str_and_repr()

Author: janosh

pymatgen/core/composition.py

@@ -151,13 +151,7 @@ def __iter__(self) -> Iterator[Species | Element | DummySpecies]:
     def __contains__(self, key) -> bool:
         try:
             sp = get_el_sp(key)
-            # First check for species
-            if sp in self._data:
-                return True
-            # If not found, check for parent element (if it's a species)
-            if isinstance(sp, Species):
-                return sp.element in self._data
-            return False
+            return sp in self._data
         except ValueError as exc:
             raise TypeError(f"Invalid {key=} for Composition") from exc
 
@@ -405,14 +399,14 @@ def get_integer_formula_and_factor(
             Li0.5O0.25 returns (Li2O, 0.25). O0.25 returns (O2, 0.125)
         """
         el_amt = self.get_el_amt_dict()
-        g = gcd_float(list(el_amt.values()), 1 / max_denominator)
+        gcd = gcd_float(list(el_amt.values()), 1 / max_denominator)
 
-        d = {k: round(v / g) for k, v in el_amt.items()}
-        formula, factor = reduce_formula(d, iupac_ordering=iupac_ordering)
+        dct = {k: round(v / gcd) for k, v in el_amt.items()}
+        formula, factor = reduce_formula(dct, iupac_ordering=iupac_ordering)
         if formula in Composition.special_formulas:
             formula = Composition.special_formulas[formula]
             factor /= 2
-        return formula, factor * g
+        return formula, factor * gcd
 
     @property
     def reduced_formula(self) -> str:
@@ -434,8 +428,8 @@ def hill_formula(self) -> str:
         no carbon, all the elements, including hydrogen, are listed
         alphabetically.
         """
-        c = self.element_composition
-        elements = sorted(el.symbol for el in c)
+        elem_comp = self.element_composition
+        elements = sorted(el.symbol for el in elem_comp)
         hill_elements = []
         if "C" in elements:
             hill_elements.append("C")
@@ -445,7 +439,7 @@ def hill_formula(self) -> str:
                 elements.remove("H")
         hill_elements += elements
 
-        formula = [f"{el}{formula_double_format(c[el]) if c[el] != 1 else ''}" for el in hill_elements]
+        formula = [f"{el}{formula_double_format(elem_comp[el]) if elem_comp[el] != 1 else ''}" for el in hill_elements]
         return " ".join(formula)
 
     @property
@@ -621,8 +615,10 @@ def valid(self) -> bool:
         """
         return not any(isinstance(el, DummySpecies) for el in self.elements)
 
-    def __repr__(self) -> str:
-        return "Comp: " + self.formula
+    def __repr__(self):
+        formula = " ".join(f"{k}{':' if hasattr(k, 'oxi_state') else ''}{v:g}" for k, v in self.items())
+        cls_name = type(self).__name__
+        return f"{cls_name}({formula!r})"
 
     @classmethod
     def from_dict(cls, d) -> Composition:
@@ -1299,9 +1295,3 @@ def __repr__(self):
 
 class CompositionError(Exception):
     """Exception class for composition errors."""
-
-
-if __name__ == "__main__":
-    import doctest
-
-    doctest.testmod()

pymatgen/core/tests/test_composition.py

@@ -89,6 +89,18 @@ def test_init(self):
         c = Composition({"S": Composition.amount_tolerance / 2})
         assert len(c.elements) == 0
 
+    def test_str_and_repr(self):
+        test_cases = [
+            ({"Li+": 2, "Mn3+": 2, "O2-": 4}, {"str": "Li+2 Mn3+2 O2-4", "repr": "Composition('Li+:2 Mn3+:2 O2-:4')"}),
+            ("H2O", {"str": "H2 O1", "repr": "Composition('H2 O1')"}),
+            ({"Fe3+": 2, "O2-": 3}, {"str": "Fe3+2 O2-3", "repr": "Composition('Fe3+:2 O2-:3')"}),
+            ("C6H6", {"str": "C6 H6", "repr": "Composition('C6 H6')"}),
+        ]
+
+        for comp, expected in test_cases:
+            assert str(Composition(comp)) == expected["str"]
+            assert repr(Composition(comp)) == expected["repr"]
+
     def test_average_electroneg(self):
         electro_negs = (2.7224999999999997, 2.4160000000000004, 2.5485714285714285, 2.21, 2.718, 3.08, 1.21, 2.43)
         for elem, val in zip(self.comps, electro_negs):

pymatgen/core/units.py

@@ -789,9 +789,9 @@ def get_mass():
 
     """
 
-    def wrap(f):
+    def wrap(func):
         def wrapped_f(*args, **kwargs):
-            val = f(*args, **kwargs)
+            val = func(*args, **kwargs)
             unit_type = _UNAME2UTYPE[unit]
 
             if isinstance(val, (FloatWithUnit, ArrayWithUnit)):
@@ -816,9 +816,3 @@ def wrapped_f(*args, **kwargs):
         return wrapped_f
 
     return wrap
-
-
-if __name__ == "__main__":
-    import doctest
-
-    doctest.testmod()