Prefer covariant types for function parameters

Closes #80

Author: michaelosthege

mcbackend/adapters/pymc.py

@@ -5,7 +5,7 @@
 """
 import base64
 import pickle
-from typing import Dict, List, Optional, Sequence, Tuple
+from typing import List, Mapping, Optional, Sequence, Tuple, Union
 
 import hagelkorn
 import numpy
@@ -59,8 +59,8 @@ def __init__(
         *,
         from_trace: BaseTrace,
         length: int,
-        draws: Dict[str, numpy.ndarray],
-        stats: Sequence[Dict[str, numpy.ndarray]],
+        draws: Mapping[str, numpy.ndarray],
+        stats: Sequence[Mapping[str, numpy.ndarray]],
     ):
         self._length = length
         self._draws = draws
@@ -105,7 +105,7 @@ def __init__(  # pylint: disable=W0622
         self,
         backend: Backend,
         *,
-        name: str = None,
+        name: Optional[str] = None,
         model=None,
         vars=None,
         test_point=None,
@@ -129,7 +129,7 @@ def setup(
         self,
         draws: int,
         chain: int,
-        sampler_vars: Optional[List[Dict[str, numpy.dtype]]] = None,
+        sampler_vars: Optional[Sequence[Mapping[str, Union[type, numpy.dtype]]]] = None,
     ) -> None:
         super().setup(draws, chain, sampler_vars)
         self.chain = chain

mcbackend/backends/clickhouse.py

@@ -5,7 +5,7 @@
 import logging
 import time
 from datetime import datetime, timezone
-from typing import Any, Callable, Dict, List, Optional, Sequence, Tuple
+from typing import Any, Callable, Dict, List, Mapping, Optional, Sequence, Tuple
 
 import clickhouse_driver
 import numpy
@@ -165,7 +165,7 @@ def __init__(
         super().__init__(cmeta, rmeta)
 
     def append(
-        self, draw: Dict[str, numpy.ndarray], stats: Optional[Dict[str, numpy.ndarray]] = None
+        self, draw: Mapping[str, numpy.ndarray], stats: Optional[Mapping[str, numpy.ndarray]] = None
     ):
         stat = {f"__stat_{sname}": svals for sname, svals in (stats or {}).items()}
         params: Dict[str, numpy.ndarray] = {**draw, **stat}

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, List, Optional, Sequence, Tuple
+from typing import Dict, List, Mapping, Optional, Sequence, Tuple
 
 import numpy
 
@@ -12,8 +12,8 @@
 
 def grow_append(
     storage_dict: Dict[str, numpy.ndarray],
-    values: Dict[str, numpy.ndarray],
-    rigid: Dict[str, bool],
+    values: Mapping[str, numpy.ndarray],
+    rigid: Mapping[str, bool],
     draw_idx: int,
 ):
     """Writes values into storage arrays, growing them if needed."""
@@ -76,7 +76,7 @@ def __init__(self, cmeta: ChainMeta, rmeta: RunMeta, *, preallocate: int) -> Non
         super().__init__(cmeta, rmeta)
 
     def append(
-        self, draw: Dict[str, numpy.ndarray], stats: Optional[Dict[str, numpy.ndarray]] = None
+        self, draw: Mapping[str, numpy.ndarray], stats: Optional[Mapping[str, numpy.ndarray]] = None
     ):
         grow_append(self._samples, draw, self._var_is_rigid, self._draw_idx)
         if stats:

mcbackend/core.py

@@ -3,7 +3,16 @@
 """
 import collections
 import logging
-from typing import TYPE_CHECKING, Dict, List, Optional, Sequence, Sized, TypeVar
+from typing import (
+    TYPE_CHECKING,
+    Dict,
+    List,
+    Mapping,
+    Optional,
+    Sequence,
+    Sized,
+    TypeVar,
+)
 
 import numpy
 
@@ -58,7 +67,7 @@ def __init__(self, cmeta: ChainMeta, rmeta: RunMeta) -> None:
         super().__init__()
 
     def append(
-        self, draw: Dict[str, numpy.ndarray], stats: Optional[Dict[str, numpy.ndarray]] = None
+        self, draw: Mapping[str, numpy.ndarray], stats: Optional[Mapping[str, numpy.ndarray]] = None
     ):
         """Appends an iteration to the chain.