Run mypy in pre-commit for continuous improvement of type hints (pymc-devs#5549)

* Remove numpy Tester instance from public API

It was making mypy follow into the tests module regardles of its config.
See https://stackoverflow.com/a/70367929

* Move all mypy config to mypy.init

Because typing_copilot can't deal with the other formats.

* Fix typing problems and add pre-commit mypy step

Author: michaelosthege

.pre-commit-config.yaml

@@ -13,6 +13,24 @@ repos:
     -   id: requirements-txt-fixer
         exclude: ^requirements-dev\.txt$
     -   id: trailing-whitespace
+- repo: https://github.com/pre-commit/mirrors-mypy
+  rev: v0.931
+  hooks:
+    - id: mypy
+      name: Run static type checks
+      language: python
+      entry: python ./scripts/run_mypy.py --verbose
+      additional_dependencies:
+      - pandas
+      - types-cachetools
+      - types-filelock
+      - types-setuptools
+      - arviz
+      - aesara==2.4.0
+      - aeppl==0.0.26
+      always_run: true
+      require_serial: true
+      pass_filenames: false
 - repo: https://github.com/PyCQA/isort
   rev: 5.10.1
   hooks:

mypy.ini

@@ -0,0 +1,12 @@
+# Autogenerated by typing_copilot v0.6.0
+[mypy]
+no_implicit_optional = False
+strict_optional = True
+warn_redundant_casts = False
+check_untyped_defs = False
+disallow_untyped_calls = False
+disallow_incomplete_defs = False
+disallow_untyped_defs = False
+disallow_untyped_decorators = False
+ignore_missing_imports = True
+warn_unused_ignores = False

pymc/__init__.py

@@ -75,7 +75,6 @@ def __set_compiler_flags():
 from pymc.smc import *
 from pymc.stats import *
 from pymc.step_methods import *
-from pymc.tests import test
 from pymc.tuning import *
 from pymc.variational import *
 from pymc.vartypes import *

pymc/backends/arviz.py

@@ -82,12 +82,9 @@ class _DefaultTrace:
         `insert()` method
     """
 
-    trace_dict: Dict[str, np.ndarray] = {}
-    _len: Optional[int] = None
-
     def __init__(self, samples: int):
-        self._len = samples
-        self.trace_dict = {}
+        self._len: int = samples
+        self.trace_dict: Dict[str, np.ndarray] = {}
 
     def insert(self, k: str, v, idx: int):
         """
@@ -180,10 +177,10 @@ def __init__(
                 " one of trace, prior, posterior_predictive or predictions."
             )
 
-        self.coords = {**self.model.coords, **(coords or {})}
+        untyped_coords = {**self.model.coords, **(coords or {})}
         self.coords = {
             cname: np.array(cvals) if isinstance(cvals, tuple) else cvals
-            for cname, cvals in self.coords.items()
+            for cname, cvals in untyped_coords.items()
             if cvals is not None
         }
 
@@ -639,7 +636,7 @@ def predictions_to_inference_data(
     """
     if inplace and not idata_orig:
         raise ValueError(
-            "Do not pass True for inplace unless passing" "an existing InferenceData as idata_orig"
+            "Do not pass True for inplace unless passing an existing InferenceData as idata_orig"
         )
     converter = InferenceDataConverter(
         trace=posterior_trace,
@@ -650,6 +647,7 @@ def predictions_to_inference_data(
         log_likelihood=False,
     )
     if hasattr(idata_orig, "posterior"):
+        assert idata_orig is not None
         converter.nchains = idata_orig.posterior.dims["chain"]
         converter.ndraws = idata_orig.posterior.dims["draw"]
     else:

pymc/blocking.py

@@ -35,6 +35,19 @@
 RaveledVars = collections.namedtuple("RaveledVars", "data, point_map_info")
 
 
+class Compose:
+    """
+    Compose two functions in a pickleable way
+    """
+
+    def __init__(self, fa: Callable[[PointType], T], fb: Callable[[RaveledVars], PointType]):
+        self.fa = fa
+        self.fb = fb
+
+    def __call__(self, x: RaveledVars) -> T:
+        return self.fa(self.fb(x))
+
+
 class DictToArrayBijection:
     """Map between a `dict`s of variables to an array space.
 
@@ -86,7 +99,9 @@ def rmap(
         return result
 
     @classmethod
-    def mapf(cls, f: Callable[[PointType], T], start_point: Optional[PointType] = None) -> T:
+    def mapf(
+        cls, f: Callable[[PointType], T], start_point: Optional[PointType] = None
+    ) -> Callable[[RaveledVars], T]:
         """Create a callable that first maps back to ``dict`` inputs and then applies a function.
 
         function f: DictSpace -> T to ArraySpace -> T
@@ -100,16 +115,3 @@ def mapf(cls, f: Callable[[PointType], T], start_point: Optional[PointType] = No
         f: array -> T
         """
         return Compose(f, partial(cls.rmap, start_point=start_point))
-
-
-class Compose:
-    """
-    Compose two functions in a pickleable way
-    """
-
-    def __init__(self, fa, fb):
-        self.fa = fa
-        self.fb = fb
-
-    def __call__(self, x):
-        return self.fa(self.fb(x))

pymc/data.py

@@ -20,7 +20,7 @@
 import warnings
 
 from copy import copy
-from typing import Any, Dict, List, Optional, Sequence, Union
+from typing import Any, Dict, List, Optional, Sequence, Union, cast
 
 import aesara
 import aesara.tensor as at
@@ -518,14 +518,15 @@ def ConstantData(
     Registers the ``value`` as a :class:`~aesara.tensor.TensorConstant` with the model.
     For more information, please reference :class:`pymc.Data`.
     """
-    return Data(
+    var = Data(
         name,
         value,
         dims=dims,
         export_index_as_coords=export_index_as_coords,
         mutable=False,
         **kwargs,
     )
+    return cast(TensorConstant, var)
 
 
 def MutableData(
@@ -541,14 +542,15 @@ def MutableData(
     Registers the ``value`` as a :class:`~aesara.compile.sharedvalue.SharedVariable`
     with the model. For more information, please reference :class:`pymc.Data`.
     """
-    return Data(
+    var = Data(
         name,
         value,
         dims=dims,
         export_index_as_coords=export_index_as_coords,
         mutable=True,
         **kwargs,
     )
+    return cast(SharedVariable, var)
 
 
 def Data(
@@ -626,9 +628,8 @@ def Data(
         value = np.array(value)
 
     # Add data container to the named variables of the model.
-    try:
-        model = pm.Model.get_context()
-    except TypeError:
+    model = pm.Model.get_context(error_if_none=False)
+    if model is None:
         raise TypeError(
             "No model on context stack, which is needed to instantiate a data container. "
             "Add variable inside a 'with model:' block."

pymc/distributions/continuous.py

@@ -158,7 +158,7 @@ class BoundedContinuous(Continuous):
     """Base class for bounded continuous distributions"""
 
     # Indices of the arguments that define the lower and upper bounds of the distribution
-    bound_args_indices = None
+    bound_args_indices: Optional[List[int]] = None
 
     def __new__(cls, *args, **kwargs):
         transform = kwargs.get("transform", UNSET)

pymc/distributions/distribution.py

@@ -186,7 +186,7 @@ class Distribution(metaclass=DistributionMeta):
     """Statistical distribution"""
 
     rv_class = None
-    rv_op = None
+    rv_op: RandomVariable = None
 
     def __new__(
         cls,

pymc/distributions/shape_utils.py

@@ -422,18 +422,18 @@ class ellipsis(Enum):
     ellipsis = type(Ellipsis)
 
 # User-provided can be lazily specified as scalars
-Shape: TypeAlias = Union[int, TensorVariable, Sequence[Union[int, TensorVariable, ellipsis]]]
+Shape: TypeAlias = Union[int, TensorVariable, Sequence[Union[int, Variable, ellipsis]]]
 Dims: TypeAlias = Union[str, Sequence[Optional[Union[str, ellipsis]]]]
-Size: TypeAlias = Union[int, TensorVariable, Sequence[Union[int, TensorVariable]]]
+Size: TypeAlias = Union[int, TensorVariable, Sequence[Union[int, Variable]]]
 
 # After conversion to vectors
-WeakShape: TypeAlias = Union[TensorVariable, Tuple[Union[int, TensorVariable, ellipsis], ...]]
+WeakShape: TypeAlias = Union[TensorVariable, Tuple[Union[int, Variable, ellipsis], ...]]
 WeakDims: TypeAlias = Tuple[Optional[Union[str, ellipsis]], ...]
 
 # After Ellipsis were substituted
-StrongShape: TypeAlias = Union[TensorVariable, Tuple[Union[int, TensorVariable], ...]]
+StrongShape: TypeAlias = Union[TensorVariable, Tuple[Union[int, Variable], ...]]
 StrongDims: TypeAlias = Sequence[Optional[str]]
-StrongSize: TypeAlias = Union[TensorVariable, Tuple[Union[int, TensorVariable], ...]]
+StrongSize: TypeAlias = Union[TensorVariable, Tuple[Union[int, Variable], ...]]
 
 
 def convert_dims(dims: Optional[Dims]) -> Optional[WeakDims]:
@@ -593,9 +593,9 @@ def find_size(
         Number of support dimensions
     """
 
-    ndim_expected = None
-    ndim_batch = None
-    create_size = None
+    ndim_expected: Optional[int] = None
+    ndim_batch: Optional[int] = None
+    create_size: Optional[StrongSize] = None
 
     if shape is not None:
         if Ellipsis in shape:

pymc/initial_point.py

@@ -68,7 +68,7 @@ def filter_rvs_to_jitter(step) -> Set[TensorVariable]:
         The random variables for which jitter should be added.
     """
     # TODO: implement this
-    return {}
+    return set()
 
 
 def make_initial_point_fns_per_chain(
@@ -163,12 +163,16 @@ def find_rng_nodes(variables):
             )
         ]
 
-    overrides = convert_str_to_rv_dict(model, overrides or {})
+    sdict_overrides = convert_str_to_rv_dict(model, overrides or {})
+    initval_strats = {
+        **model.initial_values,
+        **sdict_overrides,
+    }
 
     initial_values = make_initial_point_expression(
         free_rvs=model.free_RVs,
         rvs_to_values=model.rvs_to_values,
-        initval_strategies={**model.initial_values, **(overrides or {})},
+        initval_strategies=initval_strats,
         jitter_rvs=jitter_rvs,
         default_strategy=default_strategy,
         return_transformed=return_transformed,
@@ -178,13 +182,14 @@ def find_rng_nodes(variables):
     # when calling the final seeded function
     graph = FunctionGraph(outputs=initial_values, clone=False)
     rng_nodes = find_rng_nodes(graph.outputs)
-    new_rng_nodes = []
+    new_rng_nodes: List[Union[np.random.RandomState, np.random.Generator]] = []
     for rng_node in rng_nodes:
+        rng_cls: type
         if isinstance(rng_node, at.random.var.RandomStateSharedVariable):
-            new_rng = np.random.RandomState(np.random.PCG64())
+            rng_cls = np.random.RandomState
         else:
-            new_rng = np.random.Generator(np.random.PCG64())
-        new_rng_nodes.append(aesara.shared(new_rng))
+            rng_cls = np.random.Generator
+        new_rng_nodes.append(aesara.shared(rng_cls(np.random.PCG64())))
     graph.replace_all(zip(rng_nodes, new_rng_nodes), import_missing=True)
     func = compile_pymc(inputs=[], outputs=graph.outputs, mode=aesara.compile.mode.FAST_COMPILE)
 
@@ -310,7 +315,7 @@ def make_initial_point_expression(
 
         initial_values.append(value)
 
-    all_outputs = []
+    all_outputs: List[TensorVariable] = []
     all_outputs.extend(free_rvs)
     all_outputs.extend(initial_values)
     all_outputs.extend(initial_values_transformed)

pymc/model.py

@@ -31,6 +31,7 @@
     Type,
     TypeVar,
     Union,
+    cast,
 )
 
 import aesara
@@ -552,8 +553,8 @@ def __init__(
             self.rng_seeder = rng_seeder
 
         # The sequence of model-generated RNGs
-        self.rng_seq = []
-        self._initial_values = {}
+        self.rng_seq: List[SharedVariable] = []
+        self._initial_values: Dict[TensorVariable, Optional[Union[np.ndarray, Variable, str]]] = {}
 
         if self.parent is not None:
             self.named_vars = treedict(parent=self.parent.named_vars)
@@ -721,17 +722,20 @@ def logpt(
         -------
         Logp graph(s)
         """
+        varlist: List[TensorVariable]
         if vars is None:
-            vars = self.free_RVs + self.observed_RVs + self.potentials
+            varlist = self.free_RVs + self.observed_RVs + self.potentials
         elif not isinstance(vars, (list, tuple)):
-            vars = [vars]
+            varlist = [vars]
+        else:
+            varlist = cast(List[TensorVariable], vars)
 
         # We need to separate random variables from potential terms, and remember their
         # original order so that we can merge them together in the same order at the end
         rv_values = {}
         potentials = []
         rv_order, potential_order = [], []
-        for i, var in enumerate(vars):
+        for i, var in enumerate(varlist):
             value_var = self.rvs_to_values.get(var)
             if value_var is not None:
                 rv_values[var] = value_var
@@ -756,7 +760,7 @@ def logpt(
         if potentials:
             potential_logps, _ = rvs_to_value_vars(potentials, apply_transforms=True)
 
-        logp_factors = [None] * len(vars)
+        logp_factors = [None] * len(varlist)
         for logp_order, logp in zip((rv_order + potential_order), (rv_logps + potential_logps)):
             logp_factors[logp_order] = logp
 
@@ -948,7 +952,7 @@ def coords(self) -> Dict[str, Union[Tuple, None]]:
         return self._coords
 
     @property
-    def dim_lengths(self) -> Dict[str, Tuple[Variable, ...]]:
+    def dim_lengths(self) -> Dict[str, Variable]:
         """The symbolic lengths of dimensions in the model.
 
         The values are typically instances of ``TensorVariable`` or ``ScalarSharedVariable``.

pymc/parallel_sampling.py

@@ -21,7 +21,7 @@
 import traceback
 
 from collections import namedtuple
-from typing import Dict, Sequence
+from typing import Dict, List, Sequence
 
 import cloudpickle
 import numpy as np
@@ -425,8 +425,8 @@ def __init__(
         ]
 
         self._inactive = self._samplers.copy()
-        self._finished = []
-        self._active = []
+        self._finished: List[ProcessAdapter] = []
+        self._active: List[ProcessAdapter] = []
         self._max_active = cores
 
         self._in_context = False

pymc/tests/__init__.py

@@ -11,7 +11,3 @@
 #   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 #   See the License for the specific language governing permissions and
 #   limitations under the License.
-
-from numpy.testing import Tester
-
-test = Tester().test