Add mypy

.pre-commit-config.yaml

@@ -42,6 +42,13 @@ repos:
         args: ["--add-ignore=D107,D105"]
         additional_dependencies:
           - toml
+  - repo: https://github.com/pre-commit/mirrors-mypy
+    rev: v1.4.1
+    hooks:
+      - id: mypy
+        additional_dependencies: [types-PyYAML]
+        pass_filenames: false
+        args: ["--ignore-missing-imports", "src/"]
   - repo: https://github.com/pre-commit/pre-commit
     rev: v3.3.3
     hooks:

src/eko/couplings.py

@@ -8,10 +8,11 @@
 
 """
 import logging
-from typing import Iterable, List
+from typing import Dict, Iterable, List, Tuple
 
 import numba as nb
 import numpy as np
+import numpy.typing as npt
 import scipy
 
 from . import constants, matchings
@@ -474,7 +475,7 @@ def assert_positive(name, var):
                 self.decoupled_running,
             )
         # cache
-        self.cache = {}
+        self.cache: Dict[Tuple[float, float, int, float, float], npt.NDArray] = {}
 
     @property
     def mu2_ref(self):

src/eko/evolution_operator/__init__.py

@@ -8,6 +8,7 @@
 import os
 import time
 from multiprocessing import Pool
+from typing import Dict, Tuple
 
 import numba as nb
 import numpy as np
@@ -578,6 +579,9 @@ def quad_ker_qed(
     return ker
 
 
+OPMEMBERS = Dict[Tuple[int, int], OpMember]
+
+
 class Operator(sv.ModeMixin):
     """Internal representation of a single EKO.
 
@@ -603,8 +607,8 @@ class Operator(sv.ModeMixin):
 
     log_label = "Evolution"
     # complete list of possible evolution operators labels
-    full_labels = br.full_labels
-    full_labels_qed = br.full_unified_labels
+    full_labels = list(br.full_labels)
+    full_labels_qed = list(br.full_unified_labels)
 
     def __init__(
         self, config, managers, segment: Segment, mellin_cut=5e-2, is_threshold=False
@@ -617,7 +621,7 @@ def __init__(
         # TODO make 'cut' external parameter?
         self._mellin_cut = mellin_cut
         self.is_threshold = is_threshold
-        self.op_members = {}
+        self.op_members: OPMEMBERS = {}
         self.order = tuple(config["order"])
         self.alphaem_running = self.managers["couplings"].alphaem_running
         if self.log_label == "Evolution":

src/eko/evolution_operator/grid.py

@@ -5,8 +5,7 @@
 
 """
 import logging
-from dataclasses import astuple
-from typing import Dict, List, Optional
+from typing import Any, Dict, List, Optional
 
 import numpy as np
 import numpy.typing as npt
@@ -17,9 +16,9 @@
 from ..interpolation import InterpolatorDispatcher
 from ..io.runcards import Configs, Debug
 from ..io.types import EvolutionPoint as EPoint
-from ..io.types import Order
+from ..io.types import Order, SquaredScale
 from ..matchings import Atlas, Segment, flavor_shift, is_downward_path
-from . import Operator, flavors, matching_condition, physical
+from . import OPMEMBERS, Operator, flavors, matching_condition, physical
 from .operator_matrix_element import OperatorMatrixElement
 
 logger = logging.getLogger(__name__)
@@ -61,7 +60,7 @@ def __init__(
         interpol_dispatcher: InterpolatorDispatcher,
     ):
         # check
-        config = {}
+        config: Dict[str, Any] = {}
         config["order"] = order
         config["intrinsic_range"] = intrinsic_flavors
         config["xif2"] = xif**2
@@ -95,8 +94,8 @@ def __init__(
             couplings=couplings,
             interpol_dispatcher=interpol_dispatcher,
         )
-        self._threshold_operators = {}
-        self._matching_operators = {}
+        self._threshold_operators: Dict[Segment, Operator] = {}
+        self._matching_operators: Dict[SquaredScale, OPMEMBERS] = {}
 
     def get_threshold_operators(self, path: List[Segment]) -> List[Operator]:
         """Generate the threshold operators.
@@ -118,7 +117,7 @@ def get_threshold_operators(self, path: List[Segment]) -> List[Operator]:
         is_downward = is_downward_path(path)
         shift = flavor_shift(is_downward)
         for seg in path[:-1]:
-            new_op_key = astuple(seg)
+            new_op_key = seg
             kthr = self.config["thresholds_ratios"][seg.nf - shift]
             ome = OperatorMatrixElement(
                 self.config,

src/eko/evolution_operator/operator_matrix_element.py

@@ -3,6 +3,7 @@
 import copy
 import functools
 import logging
+from typing import List, Tuple
 
 import numba as nb
 import numpy as np
@@ -208,7 +209,7 @@ class OperatorMatrixElement(Operator):
         (br.matching_hminus_pid, br.matching_hminus_pid),
     ]
     # still valid in QED since Sdelta and Vdelta matchings are diagonal
-    full_labels_qed = copy.deepcopy(full_labels)
+    full_labels_qed: List[Tuple[int, int]] = copy.deepcopy(full_labels)
 
     def __init__(self, config, managers, nf, q2, is_backward, L, is_msbar):
         super().__init__(config, managers, Segment(q2, q2, nf))

src/eko/io/inventory.py

@@ -2,15 +2,15 @@
 import base64
 from dataclasses import asdict, dataclass, field
 from pathlib import Path
-from typing import Dict, Generic, Optional, Type, TypeVar
+from typing import Dict, Generic, Literal, Optional, Type, TypeVar
 
 import yaml
 
 from .access import AccessConfigs
 from .items import Header, Operator
 
 NBYTES = 8
-ENDIANNESS = "little"
+ENDIANNESS: Literal["little", "big"] = "little"
 
 HEADER_EXT = ".yaml"
 ARRAY_EXT = [".npy", ".npz"]

src/eko/io/struct.py

@@ -373,17 +373,17 @@ def open(cls, path: os.PathLike, mode="r"):
             raise ValueError(f"Unknown file mode: {mode}")
 
         tmpdir = pathlib.Path(tempfile.mkdtemp(prefix=TEMP_PREFIX))
-        if load:
-            cls.load(path, tmpdir)
-            metadata = Metadata.load(tmpdir)
-            opened = cls(
-                **inventories(tmpdir, access),
-                metadata=metadata,
-                access=access,
-            )
-            opened.operators.sync()
-        else:
-            opened = Builder(path=tmpdir, access=access)
+        if not load:
+            return Builder(path=tmpdir, access=access)
+        # load existing instead
+        cls.load(path, tmpdir)
+        metadata = Metadata.load(tmpdir)
+        opened: EKO = cls(
+            **inventories(tmpdir, access),
+            metadata=metadata,
+            access=access,
+        )
+        opened.operators.sync()
 
         return opened
 

src/eko/msbar_masses.py

@@ -395,7 +395,7 @@ def sc(thr_masses):
     heavy_quarks = quark_names[3:]
     hq_idxs = np.arange(0, 3)
     if nf_ref > 4:
-        heavy_quarks = reversed(heavy_quarks)
+        heavy_quarks = "".join([e for e in reversed(heavy_quarks)])
         hq_idxs = reversed(hq_idxs)
 
     # loop on heavy quarks and compute the msbar masses

src/eko/quantities/couplings.py

@@ -1,7 +1,6 @@
 """Types and quantities related to theory couplings."""
 import dataclasses
 import enum
-from typing import Optional
 
 from ..io.dictlike import DictLike
 from ..io.types import FlavorsNumber, LinearScale, ReferenceRunning, Scalar
@@ -22,7 +21,7 @@ class CouplingsInfo(DictLike):
     alphaem: Coupling
     scale: LinearScale
     max_num_flavs: FlavorsNumber
-    num_flavs_ref: Optional[FlavorsNumber]
+    num_flavs_ref: FlavorsNumber
     r"""Number of active flavors at strong coupling reference scale.
 
     I.e. :math:`n_{f,\text{ref}}(\mu^2_{\text{ref}})`, formerly called

tests/eko/quantities/test_couplings.py

@@ -3,7 +3,7 @@
 
 def test_couplings_ref():
     scale = 90.0
-    d = dict(alphas=0.1, alphaem=0.01, scale=scale, max_num_flavs=6, num_flavs_ref=None)
+    d = dict(alphas=0.1, alphaem=0.01, scale=scale, max_num_flavs=6, num_flavs_ref=5)
     couplings = CouplingsInfo.from_dict(d)
     assert couplings.scale == scale
     assert not couplings.em_running

tests/eko/test_couplings.py

@@ -54,7 +54,7 @@ def test_init(self):
                 alphas=alpharef[0],
                 alphaem=alpharef[1],
                 scale=muref,
-                num_flavs_ref=None,
+                num_flavs_ref=5,
                 max_num_flavs=6,
             )
         )
@@ -152,18 +152,19 @@ def test_ref(self):
             (0, np.inf, np.inf),
             (2, 4, 175),
         ]
+        nfrefs = (3, 4, 5)
         alpharef = (0.118, 0.00781)
         muref = 91.0
-        couplings = CouplingsInfo.from_dict(
-            dict(
-                alphas=alpharef[0],
-                alphaem=alpharef[1],
-                scale=muref,
-                num_flavs_ref=None,
-                max_num_flavs=6,
+        for thresh_setup, nfref in zip(thresh_setups, nfrefs):
+            couplings = CouplingsInfo.from_dict(
+                dict(
+                    alphas=alpharef[0],
+                    alphaem=alpharef[1],
+                    scale=muref,
+                    num_flavs_ref=nfref,
+                    max_num_flavs=6,
+                )
             )
-        )
-        for thresh_setup in thresh_setups:
             for order_s in [1, 2, 3, 4]:
                 for order_em in [0, 1, 2]:
                     for evmod in CouplingEvolutionMethod:
@@ -220,9 +221,10 @@ def test_exact(self):
             (0, np.inf, np.inf),
             (2, 4, 175),
         ]
+        nfrefs = (3, 4, 5)
         alpharef = np.array([0.118, 0.00781])
         muref = 91.0
-        for thresh_setup in thresh_setups:
+        for thresh_setup, nfref in zip(thresh_setups, nfrefs):
             for qcd in range(1, 4 + 1):
                 for qed in range(2 + 1):
                     for em_running in [
@@ -235,7 +237,7 @@ def test_exact(self):
                                 alphas=alpharef[0],
                                 alphaem=alpharef[1],
                                 scale=muref,
-                                num_flavs_ref=None,
+                                num_flavs_ref=nfref,
                                 max_num_flavs=6,
                                 em_running=em_running,
                             )
@@ -298,7 +300,7 @@ def benchmark_expanded_n3lo(self):
             alphas=alpharef[0],
             alphaem=alpharef[1],
             scale=muref,
-            num_flavs_ref=None,
+            num_flavs_ref=5,
             max_num_flavs=6,
         )
         m2c = 2