DEV: Add Type Checking With mypy

.github/workflows/ci.yml

@@ -81,7 +81,7 @@ jobs:
           flag-name: run-${{matrix.python-version}}-${{matrix.os}}
           file: "${{matrix.os}}-${{matrix.python-version}}.lcov"
 
-  finish:
+  upload-coverage:
     name: "Finish Coveralls"
     needs: tests
     runs-on: ubuntu-latest
@@ -90,4 +90,35 @@ jobs:
       uses: coverallsapp/github-action@v2
       with:
         github-token: ${{ secrets.github_token }}
-        parallel-finished: true
+        parallel-finished: true
+
+  type_check:
+    name: "Type Check"
+    needs: build-iqtree
+    runs-on: ${{ matrix.os }}
+
+    strategy:
+      matrix:
+        python-version: ["3.12"]
+        os: [ubuntu-latest]
+
+    steps:
+      - uses: "actions/checkout@v4"
+        with:
+          fetch-depth: 0
+
+      - uses: "actions/setup-python@v5"
+        with:
+            python-version: "${{ matrix.python-version }}"          
+
+      - uses: actions/cache/restore@v4
+        with:
+            key: libiqtree-${{ matrix.os }}-${{ env.IQ_TREE_2_SHA }}
+            path: src/piqtree2/_libiqtree/libiqtree2.a
+            fail-on-cache-miss: true
+
+      - name: "Run Type Checking for ${{ matrix.python-version }}" 
+        run: |
+            pip install nox
+            nox -s type_check-${{ matrix.python-version }}
+

noxfile.py

@@ -13,3 +13,13 @@ def test(session):
     install_spec = "-e.[test]"
     session.install(install_spec)
     session.run("pytest", *posargs, env=env)
+
+
+@nox.session(python=[f"3.{v}" for v in _py_versions])
+def type_check(session):
+    posargs = list(session.posargs)
+    env = os.environ.copy()
+
+    install_spec = ".[typing]"
+    session.install(install_spec)
+    session.run("mypy", "src", "tests", *posargs, env=env)

pyproject.toml

@@ -39,8 +39,10 @@ classifiers = [
 Repository = "https://github.com/cogent3/piqtree2"
 
 [project.optional-dependencies]
-dev = ["cibuildwheel", "pybind11", "ruff", "scriv", "piqtree2[test]"]
+dev = ["cibuildwheel", "pybind11", "ruff", "scriv", "piqtree2[test]", "piqtree2[typing]"]
 test = ["pytest", "pytest-cov", "nox"]
+typing = ["mypy==1.13.0", "piqtree2[stubs]", "piqtree2[test]"]
+stubs = ["types-PyYAML"]
 
 [project.entry-points."cogent3.app"]
 piqtree_phylo = "piqtree2._app:piqtree_phylo"
@@ -161,3 +163,7 @@ version = "literal: src/piqtree2/__init__.py: __version__"
 skip_fragments="README.*"
 new_fragment_template="file: changelog.d/templates/new.md.j2"
 entry_title_template="file: changelog.d/templates/title.md.j2"
+
+[[tool.mypy.overrides]]
+module = ['cogent3.*', "_piqtree2"]
+ignore_missing_imports = true

src/piqtree2/model/_freq_type.py

@@ -2,8 +2,6 @@
 import functools
 from enum import Enum, unique
 
-from typing_extensions import Self
-
 
 @unique
 class FreqType(Enum):
@@ -15,7 +13,7 @@ class FreqType(Enum):
 
     @staticmethod
     @functools.cache
-    def _descriptions() -> dict[Self, str]:
+    def _descriptions() -> dict["FreqType", str]:
         return {
             FreqType.F: "Empirical state frequency observed from the data.",
             FreqType.FO: "State frequency optimized by maximum-likelihood from the data. Note that this is with letter-O and not digit-0.",

src/piqtree2/model/_model.py

@@ -48,9 +48,8 @@ def __str__(self) -> str:
         str
             The IQ-TREE representation of the mode.
         """
-        iqtree_extra_args = filter(
-            lambda x: x is not None,
-            (self.freq_type, self.rate_type),
+        iqtree_extra_args = (
+            x for x in (self.freq_type, self.rate_type) if x is not None
         )
         return "+".join(
             x.iqtree_str() for x in [self.substitution_model, *iqtree_extra_args]

src/piqtree2/model/_options.py

@@ -16,7 +16,11 @@
 
 @functools.cache
 def _make_models(model_type: type[SubstitutionModel]) -> dict[str, list[str]]:
-    data = {"Model Type": [], "Abbreviation": [], "Description": []}
+    data: dict[str, list[str]] = {
+        "Model Type": [],
+        "Abbreviation": [],
+        "Description": [],
+    }
 
     model_classes = (
         ALL_MODELS_CLASSES if model_type == SubstitutionModel else [model_type]
@@ -62,7 +66,7 @@ def available_models(model_type: str | None = None) -> _Table:
 
 def available_freq_type() -> _Table:
     """Return a table showing available freq type options."""
-    data = {"Freq Type": [], "Description": []}
+    data: dict[str, list[str]] = {"Freq Type": [], "Description": []}
 
     for freq_type in FreqType:
         data["Freq Type"].append(freq_type.value)
@@ -73,7 +77,7 @@ def available_freq_type() -> _Table:
 
 def available_rate_type() -> _Table:
     """Return a table showing available rate type options."""
-    data = {"Rate Type": [], "Description": []}
+    data: dict[str, list[str]] = {"Rate Type": [], "Description": []}
 
     for rate_type in ALL_BASE_RATE_TYPES:
         data["Rate Type"].append(rate_type.iqtree_str())

src/piqtree2/model/_substitution_model.py

@@ -3,8 +3,6 @@
 from abc import abstractmethod
 from enum import Enum, unique
 
-from typing_extensions import Self
-
 
 class SubstitutionModel(Enum):
     """Base class for substitution models."""
@@ -23,7 +21,7 @@ def model_type() -> str:
 
     @staticmethod
     @abstractmethod
-    def _descriptions() -> dict[Self, str]:
+    def _descriptions() -> dict["SubstitutionModel", str]:
         """Get the description of each model.
 
         Returns
@@ -119,7 +117,7 @@ def model_type() -> str:
 
     @staticmethod
     @functools.cache
-    def _descriptions() -> dict[Self, str]:
+    def _descriptions() -> dict[SubstitutionModel, str]:
         return {
             DnaModel.JC: "Equal substitution rates and equal base frequencies (Jukes and Cantor, 1969).",
             DnaModel.F81: "Equal rates but unequal base freq. (Felsenstein, 1981).",
@@ -232,7 +230,7 @@ def model_type() -> str:
 
     @staticmethod
     @functools.cache
-    def _descriptions() -> dict[Self, str]:
+    def _descriptions() -> dict[SubstitutionModel, str]:
         return {
             AaModel.Blosum62: "BLOcks SUbstitution Matrix (Henikoff and Henikoff, 1992). Note that BLOSUM62 is not recommended for phylogenetic analysis as it was designed mainly for sequence alignments.",
             AaModel.cpREV: "chloroplast matrix (Adachi et al., 2000).",

tests/test_iqtree/test_build_tree.py

@@ -19,8 +19,9 @@ def check_build_tree(
     four_otu: ArrayAlignment,
     dna_model: DnaModel,
     freq_type: FreqType | None = None,
-    invariant_sites: bool | None = None,
     rate_model: RateModel | None = None,
+    *,
+    invariant_sites: bool = False,
 ) -> None:
     expected = make_tree("(Human,Chimpanzee,(Rhesus,Mouse));")
 
@@ -75,14 +76,14 @@ def test_lie_build_tree(four_otu: ArrayAlignment, dna_model: DnaModel) -> None:
 def test_rate_model_build_tree(
     four_otu: ArrayAlignment,
     dna_model: DnaModel,
-    invariant_sites: bool | None,
+    invariant_sites: bool,
     rate_model: RateModel,
 ) -> None:
     check_build_tree(
         four_otu,
         dna_model,
-        invariant_sites=invariant_sites,
         rate_model=rate_model,
+        invariant_sites=invariant_sites,
     )
 
 

tests/test_iqtree/test_segmentation_fault.py

@@ -5,7 +5,7 @@
 
 from piqtree2 import TreeGenMode, build_tree, fit_tree, random_trees
 from piqtree2.exceptions import IqTreeError
-from piqtree2.model import DiscreteGammaModel, DnaModel, FreeRateModel, Model, RateModel
+from piqtree2.model import DiscreteGammaModel, DnaModel, FreeRateModel, Model
 
 
 def test_two_build_random_trees() -> None:
@@ -39,7 +39,10 @@ def test_two_fit_random_trees() -> None:
 
 @pytest.mark.parametrize("rate_model_class", [DiscreteGammaModel, FreeRateModel])
 @pytest.mark.parametrize("categories", [0, -4])
-def test_two_invalid_models(rate_model_class: type[RateModel], categories: int) -> None:
+def test_two_invalid_models(
+    rate_model_class: type[DiscreteGammaModel] | type[FreeRateModel],
+    categories: int,
+) -> None:
     """
     Calling build_tree multiple times with an invalid
     model has resulted in a Segmentation Fault.

tests/test_model/test_options.py

@@ -12,7 +12,7 @@
     [(None, None), (DnaModel, "dna"), (AaModel, "protein")],
 )
 def test_num_available_models(
-    model_class: SubstitutionModel | None,
+    model_class: type[SubstitutionModel] | None,
     model_type: str | None,
 ) -> None:
     table = available_models(model_type)

tests/test_model/test_rate_type.py

@@ -10,7 +10,7 @@
 
 def test_rate_model_uninstantiable() -> None:
     with pytest.raises(TypeError):
-        _ = RateModel()
+        _ = RateModel()  # type: ignore[abstract]
 
 
 @pytest.mark.parametrize(
@@ -90,4 +90,7 @@ def test_invalid_rate_model_type(
         TypeError,
         match=f"Unexpected type for rate_model: {type(bad_rate_model)}",
     ):
-        _ = get_rate_type(invariant_sites=invariant_sites, rate_model=bad_rate_model)
+        _ = get_rate_type(
+            invariant_sites=invariant_sites,
+            rate_model=bad_rate_model,  # type: ignore[arg-type]
+        )