Add mypy config and type hints

Some initial values were changed to tighten type information.

Author: michaelosthege

mcbackend/adapters/pymc.py

@@ -3,7 +3,7 @@
 
 The only PyMC dependency is on the ``BaseTrace`` abstract base class.
 """
-from typing import Dict, List, Sequence, Tuple
+from typing import Dict, List, Optional, Sequence, Tuple
 
 import hagelkorn
 import numpy
@@ -104,22 +104,27 @@ def __init__(  # pylint: disable=W0622
         vars=None,
         test_point=None,
     ):
-        self.chain = None
+        self.chain: int = -1
         super().__init__(name, model, vars, test_point)
         self.run_id = hagelkorn.random(digits=6)
         print(f"Backend run id: {self.run_id}")
         self._backend: Backend = backend
 
         # Sessions created from the underlying backend
-        self._run: Run = None
-        self._chain: Chain = None
-        self._stat_groups: List[List[Tuple[str, str]]] = None
+        self._run: Optional[Run] = None
+        self._chain: Optional[Chain] = None
+        self._stat_groups: List[List[Tuple[str, str]]] = []
         self._length: int = 0
 
     def __len__(self) -> int:
         return self._length
 
-    def setup(self, draws, chain, sampler_vars=None) -> None:
+    def setup(
+        self,
+        draws: int,
+        chain: int,
+        sampler_vars: Optional[List[Dict[str, numpy.dtype]]] = None,
+    ) -> None:
         super().setup(draws, chain, sampler_vars)
         self.chain = chain
 
@@ -131,7 +136,7 @@ def setup(self, draws, chain, sampler_vars=None) -> None:
                     name,
                     str(self.var_dtypes[name]),
                     list(self.var_shapes[name]),
-                    dims=list(self.model.RV_dims[name]) if name in self.model.RV_dims else None,
+                    dims=list(self.model.RV_dims[name]) if name in self.model.RV_dims else [],
                     is_deterministic=(name not in free_rv_names),
                 )
                 for name in self.varnames
@@ -196,15 +201,23 @@ def record(self, point, sampler_states=None):
         return
 
     def get_values(self, varname, burn=0, thin=1) -> numpy.ndarray:
+        if self._chain is None:
+            raise Exception("Trace setup was not completed. Call `.setup()` first.")
         return self._chain.get_draws(varname)[burn::thin]
 
     def _get_stats(self, varname, burn=0, thin=1) -> numpy.ndarray:
+        if self._chain is None:
+            raise Exception("Trace setup was not completed. Call `.setup()` first.")
         return self._chain.get_stats(varname)[burn::thin]
 
     def _get_sampler_stats(self, stat_name, sampler_idx, burn, thin):
+        if self._chain is None:
+            raise Exception("Trace setup was not completed. Call `.setup()` first.")
         return self._get_stats(f"sampler_{sampler_idx}__{stat_name}", burn, thin)
 
     def point(self, idx: int):
+        if self._chain is None:
+            raise Exception("Trace setup was not completed. Call `.setup()` first.")
         return self._chain.get_draws_at(idx, self.var_names)
 
     def as_readonly(self) -> ReadOnlyTrace:

mcbackend/backends/clickhouse.py

@@ -5,7 +5,7 @@
 import logging
 import time
 from datetime import datetime, timezone
-from typing import Callable, Dict, Optional, Sequence, Tuple
+from typing import Any, Callable, Dict, List, Optional, Sequence, Tuple
 
 import clickhouse_driver
 import numpy
@@ -80,13 +80,13 @@ def create_chain_table(client: clickhouse_driver.Client, meta: ChainMeta, rmeta:
             )
         columns.append(column_spec_for(var, is_stat=True))
     assert len(set(columns)) == len(columns), columns
-    columns = ",\n            ".join(columns)
+    cols = ",\n            ".join(columns)
 
     query = f"""
         CREATE TABLE {cid}
         (
             `_draw_idx` UInt64,
-            {columns}
+            {cols}
         )
         ENGINE TinyLog();
     """
@@ -110,8 +110,8 @@ def __init__(
         self._client = client
         # The following attributes belong to the batched insert mechanism.
         # Inserting in batches is much faster than inserting single rows.
-        self._insert_query = None
-        self._insert_queue = []
+        self._insert_query: str = ""
+        self._insert_queue: List[Dict[str, Any]] = []
         self._last_insert = time.time()
         self._insert_interval = insert_interval
         self._insert_every = insert_every
@@ -121,7 +121,7 @@ def append(
         self, draw: Dict[str, numpy.ndarray], stats: Optional[Dict[str, numpy.ndarray]] = None
     ):
         stat = {f"__stat_{sname}": svals for sname, svals in (stats or {}).items()}
-        params = {"_draw_idx": self._draw_idx, **draw, **stat}
+        params: Dict[str, Any] = {"_draw_idx": self._draw_idx, **draw, **stat}
         self._draw_idx += 1
         if not self._insert_query:
             names = ", ".join(params.keys())
@@ -186,9 +186,10 @@ def _get_rows(  # pylint: disable=W0221
 
         # The unpacking must also account for non-rigid shapes
         if is_rigid(nshape):
+            assert nshape is not None
             buffer = numpy.empty((draws, *nshape), dtype)
         else:
-            buffer = numpy.repeat(None, draws)
+            buffer = numpy.array([None] * draws)
         for d, (vals,) in enumerate(data):
             buffer[d] = numpy.asarray(vals, dtype)
         return buffer
@@ -228,23 +229,21 @@ def __init__(
         self.created_at = created_at
         # We need handles on the chains to commit their batched inserts
         # before returning them to callers of `.get_chains()`.
-        self._chains = None
+        self._chains: List[ClickHouseChain] = []
         super().__init__(meta)
 
     def init_chain(self, chain_number: int) -> ClickHouseChain:
         cmeta = ChainMeta(self.meta.rid, chain_number)
         create_chain_table(self._client, cmeta, self.meta)
         chain = ClickHouseChain(cmeta, self.meta, client=self._client_fn())
-        if self._chains is None:
-            self._chains = []
         self._chains.append(chain)
         return chain
 
-    def get_chains(self) -> Tuple[ClickHouseChain]:
+    def get_chains(self) -> Tuple[ClickHouseChain, ...]:
         # Preferably return existing handles on chains that might have
         # uncommitted inserts pending.
         if self._chains:
-            return self._chains
+            return tuple(self._chains)
 
         # Otherwise fetch existing chains from the DB.
         chains = []
@@ -274,13 +273,14 @@ def __init__(
         """
         if client is None and client_fn is None:
             raise ValueError("Either a `client` or a `client_fn` must be provided.")
-        self._client_fn = client_fn
-        self._client = client
 
         if client_fn is None:
-            self._client_fn = lambda: client
+            client_fn = lambda: client
         if client is None:
-            self._client = self._client_fn()
+            client = client_fn()
+
+        self._client_fn: Callable[[], clickhouse_driver.Client] = client_fn
+        self._client: clickhouse_driver.Client = client
 
         create_runs_table(self._client)
         super().__init__()

mcbackend/backends/numpy.py

@@ -2,7 +2,7 @@
 This backend holds draws in memory, managing them via NumPy arrays.
 """
 import math
-from typing import Dict, Optional, Sequence, Tuple
+from typing import Dict, List, Optional, Sequence, Tuple
 
 import numpy
 
@@ -26,7 +26,7 @@ def grow_append(
             if rigid[vn]:
                 extension = numpy.empty((ngrow,) + numpy.shape(v))
             else:
-                extension = numpy.repeat(None, ngrow)
+                extension = numpy.array([None] * ngrow)
             storage_dict[vn] = numpy.concatenate((target, extension), axis=0)
             target = storage_dict[vn]
         target[draw_idx] = v
@@ -53,10 +53,10 @@ def __init__(self, cmeta: ChainMeta, rmeta: RunMeta, *, preallocate: int) -> Non
             where the correct amount of memory cannot be pre-allocated.
             In these cases, and when ``preallocate == 0`` object arrays are used.
         """
-        self._var_is_rigid = {}
-        self._samples = {}
-        self._stat_is_rigid = {}
-        self._stats = {}
+        self._var_is_rigid: Dict[str, bool] = {}
+        self._samples: Dict[str, numpy.ndarray] = {}
+        self._stat_is_rigid: Dict[str, bool] = {}
+        self._stats: Dict[str, numpy.ndarray] = {}
         self._draw_idx = 0
 
         # Create storage ndarrays for each model variable and sampler stat.
@@ -71,7 +71,7 @@ def __init__(self, cmeta: ChainMeta, rmeta: RunMeta, *, preallocate: int) -> Non
                     reserve = (preallocate, *var.shape)
                     target_dict[var.name] = numpy.empty(reserve, var.dtype)
                 else:
-                    target_dict[var.name] = numpy.repeat(None, preallocate)
+                    target_dict[var.name] = numpy.array([None] * preallocate)
 
         super().__init__(cmeta, rmeta)
 
@@ -105,7 +105,7 @@ class NumPyRun(Run):
 
     def __init__(self, meta: RunMeta, *, preallocate: int) -> None:
         self._settings = dict(preallocate=preallocate)
-        self._chains = []
+        self._chains: List[NumPyChain] = []
         super().__init__(meta)
 
     def init_chain(self, chain_number: int) -> NumPyChain:
@@ -114,7 +114,7 @@ def init_chain(self, chain_number: int) -> NumPyChain:
         self._chains.append(chain)
         return chain
 
-    def get_chains(self) -> Tuple[Chain]:
+    def get_chains(self) -> Tuple[NumPyChain, ...]:
         return tuple(self._chains)
 
 

mcbackend/core.py

@@ -3,17 +3,22 @@
 """
 import collections
 import logging
-from typing import Dict, Optional, Sequence, Sized, TypeVar
+from typing import TYPE_CHECKING, Dict, List, Optional, Sequence, Sized, TypeVar
 
 import numpy
 
 from .meta import ChainMeta, RunMeta, Variable
 from .npproto.utils import ndarray_to_numpy
 
+InferenceData = TypeVar("InferenceData")
 try:
-    from arviz import InferenceData, from_dict
+    from arviz import from_dict
+
+    if not TYPE_CHECKING:
+        from arviz import InferenceData
+    _HAS_ARVIZ = True
 except ModuleNotFoundError:
-    InferenceData = TypeVar("InferenceData")
+    _HAS_ARVIZ = False
 
 Shape = Sequence[int]
 _log = logging.getLogger(__file__)
@@ -31,7 +36,9 @@ def is_rigid(nshape: Optional[Shape]):
         - ``[2, 0]`` indicates a matrix with 2 rows and dynamic number of columns (rigid: False).
         - ``None`` indicates dynamic dimensionality (rigid: False).
     """
-    if nshape is None or any(s == 0 for s in nshape):
+    if nshape is None:
+        return False
+    if any(s == 0 for s in nshape):
         return False
     return True
 
@@ -127,14 +134,14 @@ def coords(self) -> Dict[str, numpy.ndarray]:
         return {coord.name: ndarray_to_numpy(coord.values) for coord in self.meta.coordinates}
 
     @property
-    def dims(self) -> Dict[str, Sequence[str]]:
+    def dims(self) -> Dict[str, List[str]]:
         dims = {}
         for var in self.meta.variables:
             if len(var.dims) == len(var.shape) and not var.undefined_ndim:
-                dims[var.name] = var.dims
+                dims[var.name] = list(var.dims)
         for dvar in self.meta.data:
             if len(dvar.dims) > 0:
-                dims[dvar.name] = dvar.dims
+                dims[dvar.name] = list(dvar.dims)
         return dims
 
     @property
@@ -158,7 +165,7 @@ def to_inferencedata(self, **kwargs) -> InferenceData:
         idata : arviz.InferenceData
             Samples and metadata of this inference run.
         """
-        if isinstance(InferenceData, TypeVar):
+        if not _HAS_ARVIZ:
             raise ModuleNotFoundError("ArviZ is not installed.")
 
         variables = self.meta.variables

mcbackend/npproto/utils.py

@@ -22,6 +22,7 @@ def ndarray_from_numpy(arr: numpy.ndarray) -> Ndarray:
 
 
 def ndarray_to_numpy(nda: Ndarray) -> numpy.ndarray:
+    arr: numpy.ndarray
     if "datetime64" in nda.dtype:
         # Backwards conversion: The data was stored as int64.
         arr = numpy.ndarray(

mcbackend/test_adapter_pymc.py

@@ -139,3 +139,21 @@ def wrapper(meta: RunMeta):
             ndarray_to_numpy(obs.value), simple_model["obs"].get_value()
         )
         pass
+
+    def test_uninitialized_exceptions(self, simple_model):
+        backend = ClickHouseBackend(self._client)
+        with simple_model:
+            trace = TraceBackend(backend)
+
+        with pytest.raises(Exception, match="setup was not completed"):
+            trace.get_values("scalar")
+
+        with pytest.raises(Exception, match="setup was not completed"):
+            trace._get_stats("doesntmatter")
+
+        with pytest.raises(Exception, match="setup was not completed"):
+            trace._get_sampler_stats("doesntmatter", 0, 0, 1)
+
+        with pytest.raises(Exception, match="setup was not completed"):
+            trace.point(0)
+        pass

mcbackend/test_npproto.py

@@ -16,6 +16,10 @@ class TestUtils:
             numpy.array(5),
             numpy.array(["hello", "world"]),
             numpy.array([datetime(2020, 3, 4, 5, 6, 7, 8), datetime(2020, 3, 4, 5, 6, 7, 9)]),
+            numpy.array(
+                [datetime(2020, 3, 4, 5, 6, 7, 8), datetime(2020, 3, 4, 5, 6, 7, 9)],
+                dtype="datetime64",
+            ),
             numpy.array([(1, 2), (3, 2, 1)], dtype=object),
         ],
     )

pyproject.toml

@@ -3,3 +3,9 @@ xfail_strict=true
 
 [tool.black]
 line-length = 100
+
+[tool.mypy]
+exclude = [
+  "^mcbackend/test_*",
+]
+ignore_missing_imports = true