Skip to content

Commit

Permalink
MAINT: have qn_domains-keys in Node/EdgeSettings be typed (#292)
Browse files Browse the repository at this point in the history
* now compatible with python3.12
  • Loading branch information
grayson-helmholz authored Oct 18, 2024
1 parent cebe81b commit f610cf9
Show file tree
Hide file tree
Showing 3 changed files with 59 additions and 2 deletions.
23 changes: 23 additions & 0 deletions docs/conf.py
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,7 @@

import os
import sys
import typing

from sphinx_api_relink.helpers import (
get_branch_name,
Expand All @@ -11,9 +12,20 @@
set_intersphinx_version_remapping,
)

from qrules.quantum_numbers import EdgeQuantumNumbers, NodeQuantumNumbers

sys.path.insert(0, os.path.abspath("."))
from _extend_docstrings import extend_docstrings # noqa: PLC2701


def pick_newtype_attrs(some_type: type) -> list:
return [
attr
for attr in dir(some_type)
if type(getattr(some_type, attr)) is typing.NewType
]


extend_docstrings()
set_intersphinx_version_remapping({
"ipython": {
Expand Down Expand Up @@ -261,6 +273,16 @@
nb_execution_show_tb = True
nb_execution_timeout = -1
nb_output_stderr = "remove"


nitpick_temp_names = [
*pick_newtype_attrs(EdgeQuantumNumbers),
*pick_newtype_attrs(NodeQuantumNumbers),
]
nitpick_temp_patterns = [
(r"py:(class|obj)", r"qrules\.quantum_numbers\." + name)
for name in nitpick_temp_names
]
nitpick_ignore_regex = [
(r"py:(class|obj)", "json.encoder.JSONEncoder"),
(r"py:(class|obj)", r"qrules\.topology\.EdgeType"),
Expand All @@ -269,6 +291,7 @@
(r"py:(class|obj)", r"qrules\.topology\.NewNodeType"),
(r"py:(class|obj)", r"qrules\.topology\.NodeType"),
(r"py:(class|obj)", r"qrules\.topology\.VT"),
*nitpick_temp_patterns,
]
nitpicky = True
primary_domain = "py"
Expand Down
32 changes: 32 additions & 0 deletions src/qrules/quantum_numbers.py
Original file line number Diff line number Diff line change
Expand Up @@ -127,6 +127,29 @@ class EdgeQuantumNumbers:
EdgeQuantumNumbers.g_parity,
]

# for accessing the keys of the dicts in EdgeSettings
EdgeQuantumNumberTypes = Union[
type[EdgeQuantumNumbers.pid],
type[EdgeQuantumNumbers.mass],
type[EdgeQuantumNumbers.width],
type[EdgeQuantumNumbers.spin_magnitude],
type[EdgeQuantumNumbers.spin_projection],
type[EdgeQuantumNumbers.charge],
type[EdgeQuantumNumbers.isospin_magnitude],
type[EdgeQuantumNumbers.isospin_projection],
type[EdgeQuantumNumbers.strangeness],
type[EdgeQuantumNumbers.charmness],
type[EdgeQuantumNumbers.bottomness],
type[EdgeQuantumNumbers.topness],
type[EdgeQuantumNumbers.baryon_number],
type[EdgeQuantumNumbers.electron_lepton_number],
type[EdgeQuantumNumbers.muon_lepton_number],
type[EdgeQuantumNumbers.tau_lepton_number],
type[EdgeQuantumNumbers.parity],
type[EdgeQuantumNumbers.c_parity],
type[EdgeQuantumNumbers.g_parity],
]


@frozen(init=False)
class NodeQuantumNumbers:
Expand Down Expand Up @@ -155,6 +178,15 @@ class NodeQuantumNumbers:
]
"""Type hint for quantum numbers of interaction nodes."""

# for accessing the keys of the dicts in NodeSettings
NodeQuantumNumberTypes = Union[
type[NodeQuantumNumbers.l_magnitude],
type[NodeQuantumNumbers.l_projection],
type[NodeQuantumNumbers.s_magnitude],
type[NodeQuantumNumbers.s_projection],
type[NodeQuantumNumbers.parity_prefactor],
]


def _to_optional_float(optional_float: float | None) -> float | None:
if optional_float is None:
Expand Down
6 changes: 4 additions & 2 deletions src/qrules/solving.py
Original file line number Diff line number Diff line change
Expand Up @@ -33,7 +33,9 @@
from qrules.quantum_numbers import (
EdgeQuantumNumber,
EdgeQuantumNumbers,
EdgeQuantumNumberTypes,
NodeQuantumNumber,
NodeQuantumNumberTypes,
)
from qrules.topology import MutableTransition, Topology

Expand All @@ -50,7 +52,7 @@ class EdgeSettings:

conservation_rules: set[GraphElementRule] = field(factory=set)
rule_priorities: dict[GraphElementRule, int] = field(factory=dict)
qn_domains: dict[Any, list] = field(factory=dict)
qn_domains: dict[EdgeQuantumNumberTypes, list] = field(factory=dict)


@implement_pretty_repr
Expand All @@ -70,7 +72,7 @@ class NodeSettings:

conservation_rules: set[Rule] = field(factory=set)
rule_priorities: dict[Rule, int] = field(factory=dict)
qn_domains: dict[Any, list] = field(factory=dict)
qn_domains: dict[NodeQuantumNumberTypes, list] = field(factory=dict)
interaction_strength: float = 1.0


Expand Down

0 comments on commit f610cf9

Please sign in to comment.