MAINT: make MyPy happy

Author: redeboer

src/qrules/_attrs.py

@@ -0,0 +1,17 @@
+from __future__ import annotations
+
+from fractions import Fraction
+from typing import SupportsFloat
+
+from qrules.quantum_numbers import Parity
+
+
+def to_fraction(value: SupportsFloat) -> Fraction:
+    float_value = float(value)
+    if float_value == -0.0:
+        float_value = 0.0
+    return Fraction(float_value)
+
+
+def to_parity(value: Parity | int) -> Parity:
+    return Parity(value)

src/qrules/conservation_rules.py

@@ -55,6 +55,7 @@
 from attrs import define, field, frozen
 from attrs.converters import optional
 
+from qrules._attrs import to_fraction, to_parity
 from qrules.quantum_numbers import EdgeQuantumNumbers as EdgeQN
 from qrules.quantum_numbers import NodeQuantumNumbers as NodeQN
 from qrules.quantum_numbers import arange
@@ -186,9 +187,9 @@ def parity_conservation(
 
 @frozen
 class HelicityParityEdgeInput:
-    parity: EdgeQN.parity = field(converter=EdgeQN.parity)
-    spin_magnitude: EdgeQN.spin_magnitude = field(converter=EdgeQN.spin_magnitude)
-    spin_projection: EdgeQN.spin_projection = field(converter=EdgeQN.spin_projection)
+    parity: EdgeQN.parity = field(converter=to_parity)
+    spin_magnitude: EdgeQN.spin_magnitude = field(converter=to_fraction)
+    spin_projection: EdgeQN.spin_projection = field(converter=to_fraction)
 
 
 def parity_conservation_helicity(
@@ -229,16 +230,18 @@ def parity_conservation_helicity(
 
 @frozen
 class CParityEdgeInput:
-    spin_magnitude: EdgeQN.spin_magnitude = field(converter=EdgeQN.spin_magnitude)
-    pid: EdgeQN.pid = field(converter=EdgeQN.pid)
-    c_parity: Optional[EdgeQN.c_parity] = field(converter=EdgeQN.c_parity, default=None)
+    spin_magnitude: EdgeQN.spin_magnitude = field(converter=to_fraction)
+    pid: EdgeQN.pid = field(converter=int)
+    c_parity: Optional[EdgeQN.c_parity] = field(
+        converter=optional(to_parity), default=None
+    )
 
 
 @frozen
 class CParityNodeInput:
     # These converters currently do not do anything, as "NewType"s do not have constructors
-    l_magnitude: NodeQN.l_magnitude = field(converter=NodeQN.l_magnitude)
-    s_magnitude: NodeQN.s_magnitude = field(converter=NodeQN.s_magnitude)
+    l_magnitude: NodeQN.l_magnitude = field(converter=to_fraction)
+    s_magnitude: NodeQN.s_magnitude = field(converter=to_fraction)
 
 
 def c_parity_conservation(
@@ -284,18 +287,18 @@ def _get_c_parity_multiparticle(
 
 @frozen
 class GParityEdgeInput:
-    isospin_magnitude: EdgeQN.isospin_magnitude = field(
-        converter=EdgeQN.isospin_magnitude
+    isospin_magnitude: EdgeQN.isospin_magnitude = field(converter=to_fraction)
+    spin_magnitude: EdgeQN.spin_magnitude = field(converter=to_fraction)
+    pid: EdgeQN.pid = field(converter=int)
+    g_parity: Optional[EdgeQN.g_parity] = field(
+        converter=optional(to_parity), default=None
     )
-    spin_magnitude: EdgeQN.spin_magnitude = field(converter=EdgeQN.spin_magnitude)
-    pid: EdgeQN.pid = field(converter=EdgeQN.pid)
-    g_parity: Optional[EdgeQN.g_parity] = field(converter=EdgeQN.g_parity, default=None)
 
 
 @frozen
 class GParityNodeInput:
-    l_magnitude: NodeQN.l_magnitude = field(converter=NodeQN.l_magnitude)
-    s_magnitude: NodeQN.s_magnitude = field(converter=NodeQN.s_magnitude)
+    l_magnitude: NodeQN.l_magnitude = field(converter=to_fraction)
+    s_magnitude: NodeQN.s_magnitude = field(converter=to_fraction)
 
 
 def g_parity_conservation(  # noqa: C901
@@ -375,9 +378,9 @@ def check_g_parity_isobar(
 
 @frozen
 class IdenticalParticleSymmetryOutEdgeInput:
-    spin_magnitude: EdgeQN.spin_magnitude = field(converter=EdgeQN.spin_magnitude)
-    spin_projection: EdgeQN.spin_projection = field(converter=EdgeQN.spin_projection)
-    pid: EdgeQN.pid = field(converter=EdgeQN.pid)
+    spin_magnitude: EdgeQN.spin_magnitude = field(converter=to_fraction)
+    spin_projection: EdgeQN.spin_projection = field(converter=to_fraction)
+    pid: EdgeQN.pid = field(converter=int)
 
 
 def identical_particle_symmetrization(
@@ -455,16 +458,16 @@ def _is_clebsch_gordan_coefficient_zero(
 
 @frozen
 class SpinNodeInput:
-    l_magnitude: NodeQN.l_magnitude = field(converter=NodeQN.l_magnitude)
-    l_projection: NodeQN.l_projection = field(converter=NodeQN.l_projection)
-    s_magnitude: NodeQN.s_magnitude = field(converter=NodeQN.s_magnitude)
-    s_projection: NodeQN.s_projection = field(converter=NodeQN.s_projection)
+    l_magnitude: NodeQN.l_magnitude = field(converter=to_fraction)
+    l_projection: NodeQN.l_projection = field(converter=to_fraction)
+    s_magnitude: NodeQN.s_magnitude = field(converter=to_fraction)
+    s_projection: NodeQN.s_projection = field(converter=to_fraction)
 
 
 @frozen
 class SpinMagnitudeNodeInput:
-    l_magnitude: NodeQN.l_magnitude = field(converter=NodeQN.l_magnitude)
-    s_magnitude: NodeQN.s_magnitude = field(converter=NodeQN.s_magnitude)
+    l_magnitude: NodeQN.l_magnitude = field(converter=to_fraction)
+    s_magnitude: NodeQN.s_magnitude = field(converter=to_fraction)
 
 
 def ls_spin_validity(spin_input: SpinNodeInput) -> bool:
@@ -585,12 +588,8 @@ def __spin_couplings(spin1: _Spin, spin2: _Spin) -> set[_Spin]:
 
 @define
 class IsoSpinEdgeInput:
-    isospin_magnitude: EdgeQN.isospin_magnitude = field(
-        converter=EdgeQN.isospin_magnitude
-    )
-    isospin_projection: EdgeQN.isospin_projection = field(
-        converter=EdgeQN.isospin_projection
-    )
+    isospin_magnitude: EdgeQN.isospin_magnitude = field(converter=to_fraction)
+    isospin_projection: EdgeQN.isospin_projection = field(converter=to_fraction)
 
 
 def _check_spin_valid(magnitude: Fraction, projection: Fraction) -> bool:

src/qrules/particle.py

@@ -19,13 +19,14 @@
 from fractions import Fraction
 from functools import total_ordering
 from math import copysign
-from typing import TYPE_CHECKING, Any, Callable, SupportsFloat
+from typing import TYPE_CHECKING, Any, Callable
 
 import attrs
 from attrs import field, frozen
 from attrs.converters import optional
 from attrs.validators import instance_of
 
+from qrules._attrs import to_fraction, to_parity
 from qrules.conservation_rules import GellMannNishijimaInput, gellmann_nishijima
 from qrules.quantum_numbers import Parity, _float_as_signed_str
 
@@ -44,13 +45,6 @@
 _LOGGER = logging.getLogger(__name__)
 
 
-def _to_fraction(value: SupportsFloat) -> Fraction:
-    float_value = float(value)
-    if float_value == -0.0:
-        float_value = 0.0
-    return Fraction(float_value)
-
-
 def _validate_fraction_for_spin(
     instance: Spin,
     attribute: Attribute,  # noqa: ARG001
@@ -82,11 +76,11 @@ class Spin:  # noqa: PLW1641
     """Safe, immutable data container for spin **with projection**."""
 
     magnitude: Fraction = field(
-        converter=_to_fraction,
+        converter=to_fraction,
         validator=_validate_fraction_for_spin,
     )
     projection: Fraction = field(
-        converter=_to_fraction,
+        converter=to_fraction,
         validator=_validate_fraction_for_spin,
     )
 
@@ -125,10 +119,6 @@ def _render_fraction(fraction: Fraction, plusminus: bool = False) -> str:
     return str(fraction)
 
 
-def _to_parity(value: Parity | int) -> Parity:
-    return Parity(int(value))
-
-
 def _to_spin(value: Spin | tuple[Fraction, Fraction] | tuple[float, float]) -> Spin:
     if isinstance(value, tuple):
         return Spin(*value)
@@ -173,9 +163,9 @@ class Particle:
     electron_lepton_number: int = field(default=0, validator=instance_of(int))
     muon_lepton_number: int = field(default=0, validator=instance_of(int))
     tau_lepton_number: int = field(default=0, validator=instance_of(int))
-    parity: Parity | None = field(converter=optional(_to_parity), default=None)
-    c_parity: Parity | None = field(converter=optional(_to_parity), default=None)
-    g_parity: Parity | None = field(converter=optional(_to_parity), default=None)
+    parity: Parity | None = field(converter=optional(to_parity), default=None)
+    c_parity: Parity | None = field(converter=optional(to_parity), default=None)
+    g_parity: Parity | None = field(converter=optional(to_parity), default=None)
 
     def __attrs_post_init__(self) -> None:
         if self.isospin is not None and not gellmann_nishijima(

src/qrules/quantum_numbers.py

@@ -20,7 +20,9 @@
     from collections.abc import Generator
 
 
-def _to_parity(value: int) -> Literal[-1, 1]:
+def _to_parity_int(value: int | Parity) -> Literal[-1, 1]:
+    if isinstance(value, Parity):
+        return value.value
     if not isinstance(value, int):
         msg = f"Parity must be an integer, not {type(value)}"
         raise TypeError(msg)
@@ -35,7 +37,7 @@ def _to_parity(value: int) -> Literal[-1, 1]:
 @total_ordering
 @frozen(eq=False, hash=True, order=False, repr=False)
 class Parity:  # noqa: PLW1641
-    value: Literal[-1, 1] = field(converter=_to_parity)
+    value: Literal[-1, 1] = field(converter=_to_parity_int)
 
     def __eq__(self, other: object) -> bool:
         if isinstance(other, Parity):

src/qrules/transition.py

@@ -16,6 +16,7 @@
 from attrs.validators import in_, instance_of
 from tqdm.auto import tqdm
 
+from qrules._attrs import to_fraction
 from qrules._implementers import implement_pretty_repr
 from qrules.combinatorics import (
     InitialFacts,
@@ -27,13 +28,7 @@
     match_external_edges,
     permutate_topology_kinematically,
 )
-from qrules.particle import (
-    Particle,
-    ParticleCollection,
-    ParticleWithSpin,
-    _to_fraction,
-    load_pdg,
-)
+from qrules.particle import Particle, ParticleCollection, ParticleWithSpin, load_pdg
 from qrules.quantum_numbers import (
     EdgeQuantumNumber,
     EdgeQuantumNumbers,
@@ -503,8 +498,8 @@ def create_edge_settings(edge_id: int) -> EdgeSettings:
             )
             _LOGGER.debug(
                 "using %s interaction order for node: %s",
-                str(interaction_types),
-                str(node_id),
+                interaction_types,
+                node_id,
             )
 
             temp_graph_settings: list[GraphSettings] = graph_settings
@@ -734,7 +729,7 @@ def _strip_spin(state_definition: Sequence[StateDefinition]) -> list[str]:
 @frozen(order=True)
 class State:
     particle: Particle = field(validator=instance_of(Particle))
-    spin_projection: Fraction = field(converter=_to_fraction)
+    spin_projection: Fraction = field(converter=to_fraction)
 
 
 StateTransition = FrozenTransition[State, InteractionProperties]

tests/unit/conservation_rules/test_helicity_conservation.py

@@ -11,11 +11,11 @@
     [
         (
             [
-                EdgeQuantumNumbers.spin_magnitude(s_magnitude),
+                EdgeQuantumNumbers.spin_magnitude(s_magnitude),  # type:ignore[arg-type]
             ],
             [
-                EdgeQuantumNumbers.spin_projection(lambda1),
-                EdgeQuantumNumbers.spin_projection(lambda2),
+                EdgeQuantumNumbers.spin_projection(lambda1),  # type:ignore[arg-type]
+                EdgeQuantumNumbers.spin_projection(lambda2),  # type:ignore[arg-type]
             ],
             abs(lambda1 - lambda2) <= s_magnitude,
         )