diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index cc39cc92..38a16df5 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -19,7 +19,7 @@ repos: - id: end-of-file-fixer - repo: https://github.com/astral-sh/ruff-pre-commit - rev: "v0.4.1" + rev: "v0.4.2" hooks: - id: ruff args: ["--fix", "--show-fixes"] @@ -28,7 +28,7 @@ repos: types_or: [python, pyi, jupyter] - repo: https://github.com/pre-commit/mirrors-mypy - rev: v1.9.0 + rev: v1.10.0 hooks: - id: mypy files: '^src/decaylanguage/(decay|dec|utils)/' diff --git a/notebooks/ExampleDecFileParsingWithLark.ipynb b/notebooks/ExampleDecFileParsingWithLark.ipynb index f09feb67..9a076fdb 100644 --- a/notebooks/ExampleDecFileParsingWithLark.ipynb +++ b/notebooks/ExampleDecFileParsingWithLark.ipynb @@ -338,8 +338,7 @@ " if tree.data == \"decay\":\n", " if tree.children[0].children[0].value in decays:\n", " print(\n", - " \"Decays of particle %s are redefined! Please check your .dec file.\"\n", - " % tree.children[0].children[0].value\n", + " f\"Decays of particle {tree.children[0].children[0].value} are redefined! Please check your .dec file.\"\n", " )\n", " decays[tree.children[0].children[0].value] = []\n", " for decay_mode in tree.find_data(\"decayline\"):\n", diff --git a/src/decaylanguage/dec/dec.py b/src/decaylanguage/dec/dec.py index 61e0caaa..639f55da 100644 --- a/src/decaylanguage/dec/dec.py +++ b/src/decaylanguage/dec/dec.py @@ -806,7 +806,7 @@ def _find_decay_modes(self, mother: str) -> tuple[Any, ...]: if get_decay_mother_name(decay_Tree) == mother: return tuple(decay_Tree.find_data("decayline")) - raise DecayNotFound("Decays of particle '%s' not found in .dec file!" % mother) + raise DecayNotFound(f"Decays of particle '{mother}' not found in .dec file!") def list_decay_modes(self, mother: str, pdg_name: bool = False) -> list[list[str]]: """ diff --git a/src/decaylanguage/decay/decay.py b/src/decaylanguage/decay/decay.py index 5b0d434f..9732ed3d 100644 --- a/src/decaylanguage/decay/decay.py +++ b/src/decaylanguage/decay/decay.py @@ -6,8 +6,9 @@ from __future__ import annotations +import typing from collections import Counter -from collections.abc import Sequence +from collections.abc import Collection, Sequence from copy import deepcopy from itertools import product from typing import TYPE_CHECKING, Any, Dict, Iterable, Iterator, List, TypedDict @@ -55,7 +56,7 @@ class DaughtersDict(CounterStr): def __init__( self, - iterable: dict[str, int] | list[str] | tuple[str] | str | None = None, + iterable: dict[str, int] | Collection[str] | str | None = None, **kwds: Any, ) -> None: """ @@ -569,6 +570,23 @@ def _build_decay_modes( decay_modes[mother] = DecayMode.from_dict(d) +T = typing.TypeVar("T") + + +def _get_modes(decay_chain: dict[str, T]) -> T: + # The list of decay modes is the first (and only) value of the dict + assert len(decay_chain.values()) == 1 + modes = list(decay_chain.values()) + return modes[0] + + +def _get_fs(decay: DecayModeDict) -> list[Any]: + fs = decay["fs"] + if isinstance(fs, list): + return fs + raise TypeError(f"Expected list, not {type(fs)}") + + def _expand_decay_modes( decay_chain: DecayChainDict, *, @@ -680,18 +698,6 @@ def _expand_decay_modes( } """ - def _get_modes(decay_chain: DecayChainDict) -> list[DecayModeDict]: - # The list of decay modes is the first (and only) value of the dict - assert len(decay_chain.values()) == 1 - modes = list(decay_chain.values()) - return modes[0] - - def _get_fs(decay: DecayModeDict) -> list[Any]: - fs = decay["fs"] - if isinstance(fs, list): - return fs - raise TypeError(f"Expected list, not {type(fs)}") - # The mother particle is the first (and only) key of the dict assert len(decay_chain.keys()) == 1 orig_mother = next(iter(decay_chain.keys())) @@ -705,16 +711,16 @@ def _get_fs(decay: DecayModeDict) -> list[Any]: # Replace dicts with strings (decay descriptors) expanded_modes = [] for mode in _get_modes(decay_chain): - fsp_options = [] + fsp_options: list[list[str]] = [] for fsp in _get_fs(mode): if isinstance(fsp, dict): - fsp_options += [_get_modes(fsp)] + fsp_options.append(_get_modes(fsp)) elif isinstance(fsp, str): - fsp_options += [[fsp]] # type: ignore[list-item] + fsp_options.append([fsp]) for expanded_mode in product(*fsp_options): # TODO: delegate descriptor-building to another function # allow for different conventions? - final_state = DaughtersDict(list(expanded_mode)).to_string() + final_state = DaughtersDict(expanded_mode).to_string() descriptor = DescriptorFormat.format_descriptor(mother, final_state, top) expanded_modes += [descriptor]