Add SpectralType class (#42)

astar_utils/__init__.py

@@ -5,3 +5,4 @@
 from .unique_list import UniqueList
 from .badges import Badge, BadgeReport
 from .loggers import get_logger, get_astar_logger
+from .spectral_types import SpectralType

astar_utils/spectral_types.py

@@ -0,0 +1,161 @@
+# -*- coding: utf-8 -*-
+"""Contains SpectralType class."""
+
+import re
+from typing import ClassVar
+from dataclasses import dataclass, field, InitVar
+
+
+@dataclass(frozen=True, slots=True)
+class SpectralType:
+    r"""Parse and store stellar spectral types.
+
+    This dataclass can understand a string constructor containing a valid
+    stellar spectral type including "OBAFGKM"-spectral class, 0-9 subclass
+    (including floats up to one decimal place) and luminosity class (roman
+    numerals I-V), which are converted to attributes. The initial spectral
+    type must match the regex "^[OBAFGKM](\d(\.\d)?)?(I{1,3}|IV|V)?$". Only the
+    main spectral class ("OBAFGKM") is mandatory, both numerical subclass and
+    luminosity class may be empty (None). Spectral subclasses (0.0-9.9) are
+    internally stored as floats, meaning an initial type of e.g. "A0V" is
+    considered identical to "A0.0V". Representations of an instance (``str``
+    and ``repr``) always show integers whenever possible, regardless of the
+    initially passed value.
+
+    Instances of this class are always "frozen", meaning they cannot be changed
+    after initial creation. It is not possible to manually assign the three
+    individual attribute, creation can only happen via the constructor string
+    described above.
+
+    One of the main features of this class (and indeed motivation for its
+    creation in the first place) is the ability to compare, and thus order,
+    instances of the class based on their two-component spectral (sub)class.
+    In this context, a "later" spectral type is considered "greater" than an
+    "earlier" one, i.e. O < B < A < F < G < K < M, which is consistent with the
+    convention of numerical subtypes, where 0 < 9 holds true. This is, however,
+    in contrast to the physical meaning of this parameter, which is correlated
+    with stellar effective temperature in reverse order, meaning T(A) > T(G)
+    and T(F0) > T(F9), by convention. On the other hand, in many visualisations
+    such as the Herzsprung-Russel diagram, it is common practice to represent
+    temperature on the x-axis in descending order, meaning a sorted list of
+    instances of this class will already have the typical order of such
+    diagrams.
+
+    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.
+
+    Attributes
+    ----------
+    spectral_class : str
+        Main spectral class (OBAFGKM).
+    spectral_subclass : str or None
+        Numerical spectral subclass (0.0-9.9).
+    luminosity_class : str or None
+        Roman numeral luminosity class (I-V).
+
+    Notes
+    -----
+    The constructor string can be supplied in both upper or lower case or a
+    mixture thereof, meaning "A0V", "a0v", "A0v" and "a0V" are all valid
+    representations of the same spectral type. The internal attributes are
+    converted to uppercase upon creation.
+
+    Examples
+    --------
+    >>> from astar_utils import SpectralType
+    >>> spt = SpectralType("A0V")
+    >>> spt.spectral_class
+    'A'
+
+    >>> spt.spectral_subclass
+    0.0
+
+    >>> spt.luminosity_class
+    'V'
+
+    >>> spts = [SpectralType(s) for s in
+    ...         ["G2", "M4.0", "B3", "F8", "K6.5",
+    ...          "A", "A0", "A9", "O8"]]
+    >>> sorted(spts)  # doctest: +NORMALIZE_WHITESPACE
+    [SpectralType('O8'),
+     SpectralType('B3'),
+     SpectralType('A0'),
+     SpectralType('A'),
+     SpectralType('A9'),
+     SpectralType('F8'),
+     SpectralType('G2'),
+     SpectralType('K6.5'),
+     SpectralType('M4')]
+
+    Note that a missing spectral subtype is considered as 5 (middle of the
+    main spectral class) in the context of sorting, as shown in the example
+    with the spectral type "A" ending up between "A0" and "A9".
+    """
+
+    spectral_class: str = field(init=False, default="")
+    spectral_subclass: float | None = field(init=False, default=None)
+    luminosity_class: str | None = field(init=False, default=None)
+    spectype: InitVar[str]
+    _cls_order: ClassVar = "OBAFGKM"  # descending Teff
+    _regex: ClassVar = re.compile(
+        r"^(?P<spec_cls>[OBAFGKM])(?P<sub_cls>\d(\.\d)?)?"
+        "(?P<lum_cls>I{1,3}|IV|V)?$", re.ASCII | re.IGNORECASE)
+
+    def __post_init__(self, spectype) -> None:
+        """Validate input and populate fields."""
+        if not (match := self._regex.fullmatch(spectype)):
+            raise ValueError(spectype)
+
+        classes = match.groupdict()
+        # Circumvent frozen as per the docs...
+        object.__setattr__(self, "spectral_class",
+                           str(classes["spec_cls"]).upper())
+
+        if classes["sub_cls"] is not None:
+            object.__setattr__(self, "spectral_subclass",
+                               float(classes["sub_cls"]))
+
+        if classes["lum_cls"] is not None:
+            object.__setattr__(self, "luminosity_class",
+                               str(classes["lum_cls"]).upper())
+
+    @property
+    def _subcls_str(self) -> str:
+        if self.spectral_subclass is None:
+            return ""
+        if self.spectral_subclass.is_integer():
+            return str(int(self.spectral_subclass))
+        return str(self.spectral_subclass)
+
+    def __repr__(self) -> str:
+        """Return repr(self)."""
+        return f"{self.__class__.__name__}('{self!s}')"
+
+    def __str__(self) -> str:
+        """Return str(self)."""
+        spectype = (f"{self.spectral_class}{self._subcls_str}"
+                    f"{self.luminosity_class or ''}")
+        return spectype
+
+    @property
+    def _comp_tuple(self) -> tuple[int, float]:
+        # if None, assume middle of spectral class
+        if self.spectral_subclass is not None:
+            sub_cls = self.spectral_subclass
+        else:
+            sub_cls = 5
+        return (self._cls_order.index(self.spectral_class), sub_cls)
+
+    def __lt__(self, other) -> bool:
+        """Return self < other."""
+        if not isinstance(other, self.__class__):
+            raise TypeError("Can only compare equal types.")
+        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._comp_tuple <= other._comp_tuple

pyproject.toml

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

tests/test_spectral_types.py

@@ -0,0 +1,145 @@
+# -*- coding: utf-8 -*-
+"""Unit tests for spectral_types.py."""
+
+import pytest
+import operator
+
+from astar_utils import SpectralType
+
+
+class TestParsesTypes:
+    @pytest.mark.parametrize("spec_cls", [*"OBAFGKM"])
+    def test_parses_valid_spec_cls(self, spec_cls):
+        spt = SpectralType(spec_cls)
+        assert spt.spectral_class == spec_cls
+        assert spt.spectral_subclass is None
+        assert spt.luminosity_class is None
+
+    def test_fails_on_invalid_spec_cls(self):
+        with pytest.raises(ValueError):
+            SpectralType("X")
+
+    @pytest.mark.parametrize(("spec_cls", "sub_cls"),
+                             [(*"A0",), (*"G2",), ("F", "3.5"), ("K", "7.5"),
+                              (*"M2",), (*"B9",), ("G", "5.0"), ("F", "0.5")])
+    def test_parses_valid_spec_and_sub_cls(self, spec_cls, sub_cls):
+        spt = SpectralType(spec_cls + sub_cls)
+        assert spt.spectral_class == spec_cls
+        assert spt.spectral_subclass == float(sub_cls)
+        assert spt.luminosity_class is None
+
+    @pytest.mark.parametrize("spec_sub_cls", ["X2", "Y3.5", "G99", "K1.2222"])
+    def test_fails_on_invalid_spec_and_sub_cls(self, spec_sub_cls):
+        with pytest.raises(ValueError):
+            SpectralType(spec_sub_cls)
+
+    @pytest.mark.parametrize(("spec_cls", "lum_cls"),
+                             [(*"AV",), (*"GI",), ("F", "III"), ("K", "IV")])
+    def test_parses_valid_spec_and_lum_cls(self, spec_cls, lum_cls):
+        spt = SpectralType(spec_cls + lum_cls)
+        assert spt.spectral_class == spec_cls
+        assert spt.spectral_subclass is None
+        assert spt.luminosity_class == lum_cls
+
+    @pytest.mark.parametrize("spec_lum_cls", ["II", "YVII", "GIIV", "KIVI"])
+    def test_fails_on_invalid_spec_and_lum_cls(self, spec_lum_cls):
+        with pytest.raises(ValueError):
+            SpectralType(spec_lum_cls)
+
+    @pytest.mark.parametrize(("spec_cls", "sub_cls", "lum_cls"),
+                             [(*"A0V",), (*"G2I",), ("F", "3.5", "II"),
+                              ("K", "7.5", "IV"), (*"M2V",), (*"B9I",),
+                              ("G", "5.0", "III"), ("F", "0.5", "V")])
+    def test_parses_valid_spec_sub_lum_cls(self, spec_cls, sub_cls, lum_cls):
+        spt = SpectralType(spec_cls + sub_cls + lum_cls)
+        assert spt.spectral_class == spec_cls
+        assert spt.spectral_subclass == float(sub_cls)
+        assert spt.luminosity_class == lum_cls
+
+    @pytest.mark.parametrize("spec_sub_lum_cls",
+                             ["Bogus", "G3IIV", "K4.5IVI", "B999IV"])
+    def test_fails_on_invalid_spec_sub_lum_cls(self, spec_sub_lum_cls):
+        with pytest.raises(ValueError):
+            SpectralType(spec_sub_lum_cls)
+
+
+class TestComparesTypes:
+    @pytest.mark.parametrize(("ssl_cls_a", "ssl_cls_b"),
+                             [("A0V", "A0V"), ("G2", "G2.0"),
+                              ("M3III", "M3.0III")])
+    def test_compares_classes_as_equal(self, ssl_cls_a, ssl_cls_b):
+        spt_a = SpectralType(ssl_cls_a)
+        spt_b = SpectralType(ssl_cls_b)
+        assert spt_a == spt_b
+
+    @pytest.mark.parametrize(("ssl_cls_a", "ssl_cls_b"),
+                             [("A0V", "A1V"), ("G2", "G2.5"),
+                              ("M3III", "M8.0III"), ("B3V", "BV")])
+    def test_compares_order_within_spec_cls(self, ssl_cls_a, ssl_cls_b):
+        spt_a = SpectralType(ssl_cls_a)
+        spt_b = SpectralType(ssl_cls_b)
+        assert spt_a < spt_b
+        assert spt_b > spt_a
+
+    @pytest.mark.parametrize(("ssl_cls_a", "ssl_cls_b"),
+                             [("A0V", "G1V"), ("G", "M"), ("OII", "GIV"),
+                              ("K3III", "M8.0III"), ("BV", "KI")])
+    def test_compares_order_across_spec_cls(self, ssl_cls_a, ssl_cls_b):
+        spt_a = SpectralType(ssl_cls_a)
+        spt_b = SpectralType(ssl_cls_b)
+        assert spt_a < spt_b
+        assert spt_b > spt_a
+
+    @pytest.mark.parametrize(("ssl_cls_a", "ssl_cls_b"),
+                             [("A0V", "G1V"), ("G", "M"), ("OII", "GIV"),
+                              ("K3III", "M8.0III"), ("BV", "KI"),
+                              ("A0V", "A1V"), ("G2", "G2.5"),
+                              ("M3III", "M8.0III"), ("B3V", "BV")])
+    def test_compares_order_with_equal(self, ssl_cls_a, ssl_cls_b):
+        spt_a = SpectralType(ssl_cls_a)
+        spt_b = SpectralType(ssl_cls_b)
+        assert spt_a <= spt_b
+        assert spt_b >= spt_a
+
+    @pytest.mark.parametrize("operation", [operator.gt, operator.lt,
+                                           operator.ge, operator.le])
+    def test_throws_on_invalid_compare(self, operation):
+        with pytest.raises(TypeError):
+            operation(SpectralType("A0V"), 42)
+
+
+class TestRepresentations:
+    @pytest.mark.parametrize(("ssl_cls", "exptcted"),
+                             [("A0V", "SpectralType('A0V')"),
+                              ("G2", "SpectralType('G2')"),
+                              ("K9.0", "SpectralType('K9')"),
+                              ("B2.5", "SpectralType('B2.5')"),
+                              ("M3.0III", "SpectralType('M3III')"),
+                              ("KII", "SpectralType('KII')"),])
+    def test_repr(self, ssl_cls, exptcted):
+        spt = SpectralType(ssl_cls)
+        assert repr(spt) == exptcted
+
+    @pytest.mark.parametrize(("ssl_cls", "exptcted"),
+                             [("A0V", "A0V"),
+                              ("G2", "G2"),
+                              ("K9.0", "K9"),
+                              ("B2.5", "B2.5"),
+                              ("M3.0III", "M3III"),
+                              ("KII", "KII"),])
+    def test_str(self, ssl_cls, exptcted):
+        spt = SpectralType(ssl_cls)
+        assert str(spt) == exptcted
+
+
+class TestUpperLowerCase:
+    def test_uppers_lower_case(self):
+        spt_a = SpectralType("A0V")
+        spt_b = SpectralType("a0v")
+        assert spt_a == spt_b
+
+    @pytest.mark.parametrize("mixcase", ["m2.5III", "M2.5iii"])
+    def test_uppers_lower_case_mixed(self, mixcase):
+        spt_a = SpectralType("M2.5III")
+        spt_b = SpectralType(mixcase)
+        assert spt_a == spt_b