ENH: specify allowed formalisms with Literal (#253)

docs/conf.py

@@ -64,8 +64,8 @@ def create_constraints_inventory() -> None:
 api_target_substitutions: dict[str, str | tuple[str, str]] = {
     "EdgeType": "typing.TypeVar",
     "GraphEdgePropertyMap": ("obj", "qrules.argument_handling.GraphEdgePropertyMap"),
-    "GraphNodePropertyMap": ("obj", "qrules.argument_handling.GraphNodePropertyMap"),
     "GraphElementProperties": ("obj", "qrules.solving.GraphElementProperties"),
+    "GraphNodePropertyMap": ("obj", "qrules.argument_handling.GraphNodePropertyMap"),
     "GraphSettings": ("obj", "qrules.solving.GraphSettings"),
     "InitialFacts": ("obj", "qrules.combinatorics.InitialFacts"),
     "NewEdgeType": "typing.TypeVar",
@@ -75,6 +75,7 @@ def create_constraints_inventory() -> None:
     "ParticleWithSpin": ("obj", "qrules.particle.ParticleWithSpin"),
     "qrules.topology.EdgeType": "typing.TypeVar",
     "qrules.topology.NodeType": "typing.TypeVar",
+    "SpinFormalism": ("obj", "qrules.transition.SpinFormalism"),
     "StateDefinition": ("obj", "qrules.combinatorics.StateDefinition"),
     "TypeAlias": "typing.TypeAlias",
     "typing_extensions.TypeAlias": "typing.TypeAlias",

pyproject.toml

@@ -33,7 +33,7 @@ dependencies = [
     "particle",
     "python-constraint",
     "tqdm >=4.24.0", # autonotebook
-    'typing-extensions; python_version <"3.10.0"', # typing.TypeAlias
+    'typing-extensions; python_version <"3.10.0"', # Literal, Protocol, TypeAlias
 ]
 description = "Rule-based particle reaction problem solver on a quantum number level"
 dynamic = ["version"]

src/qrules/__init__.py

@@ -59,6 +59,7 @@
     EdgeSettings,
     ProblemSet,
     ReactionInfo,
+    SpinFormalism,
     StateTransitionManager,
 )
 
@@ -264,7 +265,7 @@ def generate_transitions(  # noqa: PLR0917
     final_state: Sequence[StateDefinition],
     allowed_intermediate_particles: list[str] | None = None,
     allowed_interaction_types: str | Iterable[str] | None = None,
-    formalism: str = "canonical-helicity",
+    formalism: SpinFormalism = "canonical-helicity",
     particle_db: ParticleCollection | None = None,
     mass_conservation_factor: float | None = 3.0,
     max_angular_momentum: int = 2,

src/qrules/settings.py

@@ -50,6 +50,7 @@
 
 if TYPE_CHECKING:
     from qrules.particle import Particle, ParticleCollection
+    from qrules.transition import SpinFormalism
 
 __QRULES_PATH = dirname(realpath(__file__))
 ADDITIONAL_PARTICLES_DEFINITIONS_PATH: str = join(
@@ -118,7 +119,7 @@ def from_str(description: str) -> InteractionType:
 
 
 def create_interaction_settings(  # noqa: PLR0917
-    formalism: str,
+    formalism: SpinFormalism,
     particle_db: ParticleCollection,
     nbody_topology: bool = False,
     mass_conservation_factor: float | None = 3.0,

src/qrules/transition.py

@@ -14,7 +14,7 @@
 
 import attrs
 from attrs import define, field, frozen
-from attrs.validators import instance_of
+from attrs.validators import in_, instance_of
 from tqdm.auto import tqdm
 
 from qrules._implementers import implement_pretty_repr
@@ -74,6 +74,11 @@
     create_n_body_topology,
 )
 
+if sys.version_info >= (3, 8):
+    from typing import Literal
+else:
+    from typing_extensions import Literal
+
 if sys.version_info >= (3, 10):
     from typing import TypeAlias
 else:
@@ -83,6 +88,19 @@
 
 _LOGGER = logging.getLogger(__name__)
 
+SpinFormalism = Literal[
+    "helicity",
+    "canonical-helicity",
+    "canonical",
+]
+"""Name for the spin formalism to be used.
+
+The options :code:`"helicity"`, :code:`"canonical-helicity"`, and :code:`"canonical"`
+are all used for the helicity formalism, but :code:`"canonical-helicity"` and
+:code:`"canonical"` generate angular momentum and coupled spins as well on the
+interaction nodes.
+"""
+
 
 class SolvingMode(Enum):
     """Types of modes for solving."""
@@ -226,7 +244,7 @@ def __init__(  # noqa: C901, PLR0912, PLR0917
             InteractionType, tuple[EdgeSettings, NodeSettings]
         ]
         | None = None,
-        formalism: str = "helicity",
+        formalism: SpinFormalism = "helicity",
         topology_building: str = "isobar",
         solving_mode: SolvingMode = SolvingMode.FAST,
         reload_pdg: bool = False,
@@ -240,18 +258,13 @@ def __init__(  # noqa: C901, PLR0912, PLR0917
         self.__number_of_threads = NumberOfThreads.get()
         if interaction_type_settings is None:
             interaction_type_settings = {}
-        allowed_formalisms = [
-            "helicity",
-            "canonical-helicity",
-            "canonical",
-        ]
-        if formalism not in allowed_formalisms:
+        if formalism not in set(SpinFormalism.__args__):  # type: ignore[attr-defined]
             msg = (
                 f'Formalism "{formalism}" not implemented. Use one of'
-                f" {allowed_formalisms} instead."
+                f" {', '.join(SpinFormalism.__args__)} instead."  # type: ignore[attr-defined]
             )
             raise NotImplementedError(msg)
-        self.__formalism = str(formalism)
+        self.__formalism = formalism
         self.__particles = ParticleCollection()
         if particle_db is not None:
             self.__particles = particle_db
@@ -343,7 +356,7 @@ def set_allowed_intermediate_particles(
         self.__intermediate_particle_filters = selected_particles.names
 
     @property
-    def formalism(self) -> str:
+    def formalism(self) -> SpinFormalism:
         return self.__formalism
 
     def add_final_state_grouping(self, fs_group: list[str] | list[list[str]]) -> None:
@@ -744,7 +757,7 @@ class ReactionInfo:
     """Ordered collection of `StateTransition` instances."""
 
     transitions: tuple[StateTransition, ...] = field(converter=_sort_tuple)
-    formalism: str = field(validator=instance_of(str))
+    formalism: SpinFormalism = field(validator=in_(SpinFormalism.__args__))  # type: ignore[attr-defined]
 
     initial_state: FrozenDict[int, Particle] = field(init=False, repr=False, eq=False)
     final_state: FrozenDict[int, Particle] = field(init=False, repr=False, eq=False)

tests/channels/test_psi2s_to_eta_k_kstar.py

@@ -4,6 +4,7 @@
 
 import qrules
 from qrules.particle import ParticleCollection
+from qrules.transition import SpinFormalism
 
 
 @pytest.mark.parametrize("formalism", ["helicity", "canonical-helicity"])
@@ -15,7 +16,7 @@
         ["h(1)(1415)", "omega(1650)"],
     ],
 )
-def test_resonances(formalism, resonances, modified_pdg):
+def test_resonances(formalism: SpinFormalism, resonances, modified_pdg):
     reaction = qrules.generate_transitions(
         initial_state=("psi(2S)", [+1, -1]),
         final_state=["eta", "K-", "K*(892)+"],

tests/unit/conftest.py

@@ -1,4 +1,5 @@
 import logging
+from typing import TYPE_CHECKING
 
 import pytest
 from _pytest.fixtures import SubRequest
@@ -7,12 +8,15 @@
 from qrules import ReactionInfo
 from qrules.topology import Edge, Topology
 
+if TYPE_CHECKING:
+    from qrules.transition import SpinFormalism
+
 logging.basicConfig(level=logging.ERROR)
 
 
 @pytest.fixture(scope="session", params=["canonical-helicity", "helicity"])
 def reaction(request: SubRequest) -> ReactionInfo:
-    formalism: str = request.param
+    formalism: SpinFormalism = request.param
     return qrules.generate_transitions(
         initial_state=[("J/psi(1S)", [-1, 1])],
         final_state=["gamma", "pi0", "pi0"],

tests/unit/io/test_dot.py

@@ -11,7 +11,7 @@
     create_isobar_topologies,
     create_n_body_topology,
 )
-from qrules.transition import ReactionInfo
+from qrules.transition import ReactionInfo, SpinFormalism
 
 
 def test_asdot(reaction: ReactionInfo):
@@ -116,7 +116,7 @@ def test_asdot_no_label_overwriting(reaction: ReactionInfo):
     "formalism",
     ["canonical", "canonical-helicity", "helicity"],
 )
-def test_asdot_problemset(formalism: str):
+def test_asdot_problemset(formalism: SpinFormalism):
     stm = qrules.StateTransitionManager(
         initial_state=[("J/psi(1S)", [+1])],
         final_state=["gamma", "pi0", "pi0"],

tests/unit/test_settings.py

@@ -9,6 +9,7 @@
     _int_domain,
     create_interaction_settings,
 )
+from qrules.transition import SpinFormalism
 
 
 class TestInteractionType:
@@ -60,7 +61,7 @@ def test_create_interaction_settings(
     particle_database: ParticleCollection,
     interaction_type: InteractionType,
     nbody_topology: bool,
-    formalism: str,
+    formalism: SpinFormalism,
 ):
     settings = create_interaction_settings(
         formalism,