mypy init

.github/workflows/test.yml

@@ -43,6 +43,10 @@ jobs:
           echo "Success!"
           echo "Checking code style with pylint..."
           python -m pylint pymc_bart/
+      - name: Run Mypy
+        shell: bash -l {0}
+        run: |
+          python -m mypy pymc_bart
       - name: Run tests
         shell: bash -l {0}
         run: |

.pre-commit-config.yaml

@@ -13,6 +13,12 @@ repos:
         language: system
         types: [python]
         files: ^pymc_bart/
+      - id: mypy
+        name: mypy
+        entry: mypy
+        language: system
+        types: [python]
+        files: ^pymc_bart/
   - repo: https://github.com/pre-commit/pre-commit-hooks
     rev: v3.2.0
     hooks:

mypy.ini

@@ -0,0 +1,15 @@
+[mypy]
+files = pymc_bart/*.py
+plugins = numpy.typing.mypy_plugin
+
+[mypy-matplotlib.*]
+ignore_missing_imports = True
+
+[mypy-numba.*]
+ignore_missing_imports = True
+
+[mypy-pymc.*]
+ignore_missing_imports = True
+
+[mypy-scipy.*]
+ignore_missing_imports = True

pymc_bart/bart.py

@@ -15,29 +15,31 @@
 #   limitations under the License.
 
 from multiprocessing import Manager
+from typing import List, Optional, Tuple
+
 import numpy as np
+import numpy.typing as npt
+import pytensor.tensor as pt
 from pandas import DataFrame, Series
-
 from pymc.distributions.distribution import Distribution, _moment
 from pymc.logprob.abstract import _logprob
-import pytensor.tensor as pt
 from pytensor.tensor.random.op import RandomVariable
 
-
-from .utils import _sample_posterior
+from .tree import Tree
+from .utils import TensorLike, _sample_posterior
 
 __all__ = ["BART"]
 
 
 class BARTRV(RandomVariable):
     """Base class for BART."""
 
-    name = "BART"
+    name: str = "BART"
     ndim_supp = 1
-    ndims_params = [2, 1, 0, 0, 1]
-    dtype = "floatX"
-    _print_name = ("BART", "\\operatorname{BART}")
-    all_trees = None
+    ndims_params: List[int] = [2, 1, 0, 0, 1]
+    dtype: str = "floatX"
+    _print_name: Tuple[str, str] = ("BART", "\\operatorname{BART}")
+    all_trees = List[List[Tree]]
 
     def _supp_shape_from_params(self, dist_params, rep_param_idx=1, param_shapes=None):
         return dist_params[0].shape[:1]
@@ -64,29 +66,29 @@ class BART(Distribution):
 
     Parameters
     ----------
-    X : array-like
+    X : TensorLike
         The covariate matrix.
-    Y : array-like
+    Y : TensorLike
         The response vector.
     m : int
         Number of trees
     alpha : float
         Control the prior probability over the depth of the trees. Even when it can takes values in
         the interval (0, 1), it is recommended to be in the interval (0, 0.5].
-    split_prior : array-like
+    split_prior : Optional[List[float]], default None.
         Each element of split_prior should be in the [0, 1] interval and the elements should sum to
         1. Otherwise they will be normalized.
         Defaults to 0, i.e. all covariates have the same prior probability to be selected.
     """
 
     def __new__(
         cls,
-        name,
-        X,
-        Y,
-        m=50,
-        alpha=0.25,
-        split_prior=None,
+        name: str,
+        X: TensorLike,
+        Y: TensorLike,
+        m: int = 50,
+        alpha: float = 0.25,
+        split_prior: Optional[List[float]] = None,
         **kwargs,
     ):
         manager = Manager()
@@ -147,7 +149,9 @@ def get_moment(cls, rv, size, *rv_inputs):
         return mean
 
 
-def preprocess_xy(X, Y):
+def preprocess_xy(
+    X: TensorLike, Y: TensorLike
+) -> Tuple[npt.NDArray[np.float_], npt.NDArray[np.float_]]:
     if isinstance(Y, (Series, DataFrame)):
         Y = Y.to_numpy()
     if isinstance(X, (Series, DataFrame)):