Make tests pass.

Author: jesnie

check_shapes/__init__.py

@@ -451,9 +451,9 @@
 
 * Python built-in scalars: ``bool``, ``int``, ``float`` and ``str``.
 * Python built-in sequences: ``tuple`` and ``list``.
-* NumPy ``ndarray``\ s.
-* TensorFlow ``Tensor``\ s and ``Variable``\ s.
-* TensorFlow Probability ``DeferredTensor``\ s, including ``TransformedVariable`` and
+* NumPy ``ndarray``\\ s.
+* TensorFlow ``Tensor``\\ s and ``Variable``\\ s.
+* TensorFlow Probability ``DeferredTensor``\\ s, including ``TransformedVariable`` and
   :class:`gpflow.Parameter`.
 
 
@@ -490,6 +490,8 @@
 from .inheritance import inherit_check_shapes
 from .shapes import get_shape, register_get_shape
 
+__version__ = "0.1.0"
+
 __all__ = [
     "Dimension",
     "DocstringFormat",

check_shapes/accessors.py

@@ -34,7 +34,7 @@ def maybe_get_check_shapes(func: Callable[..., Any]) -> Optional[Callable[[C], C
     :returns: The ``check_shapes`` that is wrapping ``func``, and ``None`` if no ``check_shapes`` is
         not wrapping ``func``.
     """
-    return getattr(func, "__check_shapes__", None)
+    return getattr(func, "__check_shapes__", None)  # type: ignore[no-any-return]
 
 
 def get_check_shapes(func: Callable[..., Any]) -> Callable[[C], C]:

check_shapes/argument_ref.py

@@ -96,7 +96,7 @@ def __repr__(self) -> str:
 
 @dataclass(frozen=True)  # type: ignore[misc]
 class DelegatingArgumentRef(ArgumentRef):
-    """ Abstract base class for :class:`ArgumentRef`\ s that delegates to a source. """
+    """ Abstract base class for :class:`ArgumentRef`\\ s that delegates to a source. """
 
     source: ArgumentRef
 

check_shapes/base_types.py

@@ -14,7 +14,11 @@
 """
 Definitions of commonly used types.
 """
-from typing import Any, Callable, Optional, Tuple, TypeVar
+from typing import TYPE_CHECKING, Any, Callable, Optional, Tuple, TypeVar, Union
+
+import numpy as np
+
+from .type_flags import GENERIC_NP_ARRAYS, NP_TYPE_CHECKING
 
 C = TypeVar("C", bound=Callable[..., Any])
 
@@ -33,3 +37,14 @@
 
 Raise an exception if objects of that type can never have a shape.
 """
+
+if TYPE_CHECKING and (not NP_TYPE_CHECKING):  # pragma: no cover
+    AnyNDArray = Any
+else:
+    if GENERIC_NP_ARRAYS:
+        # It would be nice to use something more interesting than `Any` here, but it looks like
+        # the infrastructure in the rest of the ecosystem isn't really set up for this
+        # yet. Maybe when we get Python 3.11?
+        AnyNDArray = np.ndarray[Any, Any]  # type: ignore[misc]
+    else:
+        AnyNDArray = Union[np.ndarray]  # type: ignore[misc]

check_shapes/bool_specs.py

@@ -20,7 +20,12 @@
 from typing import Any, Callable, Mapping, Tuple, cast
 
 from .argument_ref import ArgumentRef
-from .error_contexts import ErrorContext, ObjectValueContext, ParallelContext, StackContext
+from .error_contexts import (
+    ErrorContext,
+    ObjectValueContext,
+    ParallelContext,
+    StackContext,
+)
 
 
 class ParsedBoolSpec(ABC):

check_shapes/checker.py

@@ -15,9 +15,19 @@
 Class responsible for remembering and checking shapes.
 """
 from dataclasses import dataclass, field
-from typing import Any, DefaultDict, Dict, Iterable, List, Optional, Set, Tuple, TypeVar, Union
+from typing import (
+    Any,
+    DefaultDict,
+    Dict,
+    Iterable,
+    List,
+    Optional,
+    Set,
+    Tuple,
+    TypeVar,
+    Union,
+)
 
-from ..utils import experimental
 from .base_types import Dimension, Shape
 from .config import get_enable_check_shapes
 from .error_contexts import (
@@ -227,7 +237,6 @@ class ShapeChecker:
        :dedent:
     """
 
-    @experimental
     def __init__(self) -> None:
         self._variables: DefaultDict[str, _VariableState] = DefaultDict(_VariableState)
         self._observed_shapes: List[Tuple[Shape, ParsedTensorSpec, ErrorContext]] = []

check_shapes/decorator.py

@@ -20,7 +20,6 @@
 
 import tensorflow as tf
 
-from ..utils import experimental
 from .accessors import set_check_shapes
 from .argument_ref import RESULT_TOKEN
 from .base_types import C
@@ -50,7 +49,6 @@ def null_check_shapes(func: C) -> C:
     return func
 
 
-@experimental
 def check_shapes(*specs: str, tf_decorator: bool = False) -> Callable[[C], C]:
     """
     Decorator that checks the shapes of tensor arguments.

check_shapes/error_contexts.py

@@ -22,7 +22,7 @@
 The :class:`ErrorContext` is a reusable bit of information about where/why an error occurred that
 can be written to a :class:`MessageBuilder`.
 
-:class:`ErrorContext`\ s can be composed using the :class:`StackContext` and
+:class:`ErrorContext`\\ s can be composed using the :class:`StackContext` and
 :class:`ParallelContext`.
 
 This allows reusable error messages in a consistent format.
@@ -59,7 +59,7 @@
 _UNKNOWN_LINE = "<Unknown line>"
 _DISABLED_FILE_AND_LINE = (
     f"{_UNKNOWN_FILE}:{_UNKNOWN_LINE}"
-    " (Disabled. Call gpflow.experimental.check_shapes.set_enable_function_call_precompute(True) to"
+    " (Disabled. Call check_shapes.set_enable_function_call_precompute(True) to"
     " see this. (Slow.))"
 )
 _NONE_SHAPE = "<Tensor is None or has unknown shape>"
@@ -225,7 +225,9 @@ def _split_head(context: StackContext) -> Tuple[ErrorContext, ErrorContext]:
         by_head: Dict[ErrorContext, List[ErrorContext]] = {}
         for child in flat:
             if isinstance(child, StackContext):
-                head, body = _split_head(child)
+                split_child = _split_head(child)
+                head: ErrorContext = split_child[0]
+                body: Optional[ErrorContext] = split_child[1]
             else:
                 head = child
                 body = None

check_shapes/inheritance.py

@@ -17,14 +17,12 @@
 import inspect
 from typing import Callable, Optional, cast
 
-from ..utils import experimental
 from .accessors import maybe_get_check_shapes
 from .base_types import C
 from .config import get_enable_check_shapes
 from .decorator import null_check_shapes
 
 
-@experimental
 def inherit_check_shapes(func: C) -> C:
     """
     Decorator that inherits the :func:`check_shapes` decoration from any overridden method in a

check_shapes/mypy_flags.py

@@ -0,0 +1,25 @@
+# Copyright 2022 The GPflow Contributors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# 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.
+"""
+Code for printing flags for mypy, depending on library versions.
+"""
+from .type_flags import compute_mypy_flags
+
+
+def print_mypy_flags() -> None:
+    print(compute_mypy_flags())
+
+
+if __name__ == "__main__":
+    print_mypy_flags()