Allow comparison of SpectralType directly to string (#45)

astar_utils/spectral_types.py

@@ -44,7 +44,9 @@ class SpectralType:
     In this context, the luminosity class (if any) is ignored for sorting and
     comparison (<, >, <=, >=), as it represents a second physical dimension.
     However, instances of this class may also be compared for equality (== and
-    !=), in which case all three attributes are considered.
+    !=), in which case all three attributes are considered. It is also possible
+    to compare instances directly to strings, if the string is a valid
+    construtor for this class.
 
     Attributes
     ----------
@@ -106,7 +108,7 @@ class SpectralType:
     def __post_init__(self, spectype) -> None:
         """Validate input and populate fields."""
         if not (match := self._regex.fullmatch(spectype)):
-            raise ValueError(spectype)
+            raise ValueError(f"{spectype!r} is not a valid spectral type.")
 
         classes = match.groupdict()
         # Circumvent frozen as per the docs...
@@ -178,14 +180,35 @@ def _comp_tuple(self) -> tuple[int, float]:
             sub_cls = 5
         return (self._spec_cls_idx, sub_cls)
 
+    @classmethod
+    def _comp_guard(cls, other):
+        if isinstance(other, str):
+            other = cls(other)
+        if not isinstance(other, cls):
+            raise TypeError("Can only compare equal types or valid str.")
+        return other
+
+    def __eq__(self, other) -> bool:
+        """Return self == other."""
+        other = self._comp_guard(other)
+        return self._comp_tuple == other._comp_tuple
+
     def __lt__(self, other) -> bool:
         """Return self < other."""
-        if not isinstance(other, self.__class__):
-            raise TypeError("Can only compare equal types.")
+        other = self._comp_guard(other)
         return self._comp_tuple < other._comp_tuple
 
     def __le__(self, other) -> bool:
-        """Return self < other."""
-        if not isinstance(other, self.__class__):
-            raise TypeError("Can only compare equal types.")
+        """Return self <= other."""
+        other = self._comp_guard(other)
         return self._comp_tuple <= other._comp_tuple
+
+    def __gt__(self, other) -> bool:
+        """Return self > other."""
+        other = self._comp_guard(other)
+        return self._comp_tuple > other._comp_tuple
+
+    def __ge__(self, other) -> bool:
+        """Return self >= other."""
+        other = self._comp_guard(other)
+        return self._comp_tuple >= other._comp_tuple

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "astar-utils"
-version = "0.3.1a1"
+version = "0.3.1a2"
 description = "Contains commonly-used utilities for AstarVienna's projects."
 license = "GPL-3.0-or-later"
 authors = ["Fabian Haberhauer <fabian.haberhauer@univie.ac.at>"]

tests/test_spectral_types.py

@@ -108,6 +108,35 @@ def test_throws_on_invalid_compare(self, operation):
             operation(SpectralType("A0V"), 42)
 
 
+class TestComparesStr:
+    def test_lt(self):
+        assert SpectralType("A0V") < "A7V"
+
+    def test_le(self):
+        assert SpectralType("A0V") <= "A0V"
+
+    def test_gt(self):
+        assert SpectralType("A0V") > "B7V"
+
+    def test_ge(self):
+        assert SpectralType("A0V") >= "A0V"
+
+    def test_eq(self):
+        assert SpectralType("A0V") == "A0V"
+
+    def test_ne(self):
+        assert SpectralType("A0V") != "A1V"
+
+    def test_reverse_le(self):
+        assert "A0" <= SpectralType("A0V")
+
+    def test_reverse_gt(self):
+        assert "A7" > SpectralType("A0V")
+
+    def test_reverse_ne(self):
+        assert "A1" != SpectralType("A0V")
+
+
 class TestRepresentations:
     @pytest.mark.parametrize(("ssl_cls", "exptcted"),
                              [("A0V", "SpectralType('A0V')"),