ENH: support MatrixSymbol in HelicityModel attributes (#495)

* ENH: make signature in dict sorting functions generic

Author: Zeyna777

docs/conf.py

@@ -115,6 +115,7 @@ def _get_scipy_url() -> str:
     "FourMomenta": ("obj", "ampform.kinematics.lorentz.FourMomenta"),
     "FourMomentumSymbol": ("obj", "ampform.kinematics.lorentz.FourMomentumSymbol"),
     "InteractionProperties": "qrules.quantum_numbers.InteractionProperties",
+    "K": "typing.TypeVar",
     "LatexPrinter": "sympy.printing.printer.Printer",
     "Literal[-1, 1]": "typing.Literal",
     "Literal[(-1, 1)]": "typing.Literal",
@@ -139,6 +140,7 @@ def _get_scipy_url() -> str:
     "SympyObject": "typing.TypeVar",
     "T": "typing.TypeVar",
     "Topology": "qrules.topology.Topology",
+    "V": "typing.TypeVar",
     "WignerD": "sympy.physics.quantum.spin.WignerD",
     "sympy.tensor.array.expressions.array_expressions.ArraySymbol": (
         "mod",

src/ampform/helicity/__init__.py

@@ -14,7 +14,7 @@
 from collections import OrderedDict, abc
 from fractions import Fraction
 from functools import reduce, singledispatchmethod
-from typing import TYPE_CHECKING, cast
+from typing import TYPE_CHECKING, TypeVar, cast
 
 import attrs
 import sympy as sp
@@ -74,27 +74,24 @@
 
 _LOGGER = logging.getLogger(__name__)
 
+K = TypeVar("K")
+V = TypeVar("V")
 
-def _order_component_mapping(
-    mapping: Mapping[str, sp.Expr],
-) -> OrderedDict[str, sp.Expr]:
+
+def _order_component_mapping(mapping: Mapping[str, V], /) -> OrderedDict[str, V]:
     return collections.OrderedDict([
         (key, mapping[key]) for key in sorted(mapping, key=natural_sorting)
     ])
 
 
-def _order_symbol_mapping(
-    mapping: Mapping[sp.Symbol, sp.Expr],
-) -> OrderedDict[sp.Symbol, sp.Expr]:
+def _order_symbol_mapping(mapping: Mapping[K, V], /) -> OrderedDict[K, V]:
     return collections.OrderedDict([
         (symbol, mapping[symbol])
-        for symbol in sorted(mapping, key=lambda s: natural_sorting(s.name))
+        for symbol in sorted(mapping, key=lambda s: natural_sorting(str(s)))
     ])
 
 
-def _order_amplitudes(
-    mapping: Mapping[sp.Indexed, sp.Expr],
-) -> OrderedDict[sp.Indexed, sp.Expr]:
+def _order_amplitudes(mapping: Mapping[K, V], /) -> OrderedDict[K, V]:
     return collections.OrderedDict([
         (key, mapping[key])
         for key in sorted(mapping, key=lambda a: natural_sorting(str(a)))