Static type-checking fix (#34)

* Add support for e.g. jaxtyping.Float[Union[...], ...] in py3.8

Turns out that python3.8, Union lacks the __name__ attribute.  Use repr()
in these cases.

* Fix linter.

* Remove implicit cast to bool in favor of try/except.

MANIFEST.in

@@ -1,2 +1,2 @@
-include LICENSE
+include jaxtyping/py.typed
-prune tests
+prune test

jaxtyping/__init__.py

@@ -20,14 +20,18 @@
 import typing
 
 
-if getattr(typing, "GENERATING_DOCUMENTATION", False):
+if typing.TYPE_CHECKING:
-
+    # type checkers don't know which branch below will be executed
-    class Array:
-        pass
-
-    Array.__module__ = "builtins"
-else:
     from jax.numpy import ndarray as Array
+else:
+    if getattr(typing, "GENERATING_DOCUMENTATION", False):
+
+        class Array:
+            pass
+
+        Array.__module__ = "builtins"
+    else:
+        from jax.numpy import ndarray as Array
 
 from .array_types import (
     AbstractArray,
@@ -63,4 +67,4 @@ from .import_hook import install_import_hook
 from .pytree_type import PyTree
 
 
-__version__ = "0.2.4"
+__version__ = "0.2.6"

jaxtyping/array_types.py

@@ -150,26 +150,25 @@ class _MetaAbstractArray(type):
         if cls.dtypes is not _any_dtype and dtype not in cls.dtypes:
             return False
 
-        temp_memo = not hasattr(storage, "memo_stack") or len(storage.memo_stack) == 0
+        no_temp_memo = hasattr(storage, "memo_stack") and len(storage.memo_stack) != 0
 
-        if temp_memo:
+        if no_temp_memo:
+            single_memo, variadic_memo, variadic_broadcast_memo = storage.memo_stack[-1]
+            # Make a copy so we don't mutate the original memo during the shape check.
+            single_memo = single_memo.copy()
+            variadic_memo = variadic_memo.copy()
+            variadic_broadcast_memo = variadic_broadcast_memo.copy()
+        else:
             # `isinstance` happening outside any @jaxtyped decorators, e.g. at the
             # global scope. In this case just create a temporary memo, since we're not
             # going to be comparing against any stored values anyway.
             single_memo = {}
             variadic_memo = {}
             variadic_broadcast_memo = {}
-        else:
-            single_memo, variadic_memo, variadic_broadcast_memo = storage.memo_stack[-1]
-            # Make a copy so we don't mutate the original memo during the shape check.
-            single_memo = single_memo.copy()
-            variadic_memo = variadic_memo.copy()
-            variadic_broadcast_memo = variadic_broadcast_memo.copy()
-            temp_memo = False
 
         if cls._check_shape(obj, single_memo, variadic_memo, variadic_broadcast_memo):
             # We update the memo every time we successfully pass a shape check
-            if not temp_memo:
+            if no_temp_memo:
                 storage.memo_stack[-1] = (
                     single_memo,
                     variadic_memo,
@@ -385,7 +384,12 @@ class _MetaAbstractDtype(type):
                 elem = _SymbolicDim(elem, broadcastable)
             dims.append(elem)
         if _array_name_format == "dtype_and_shape":
-            name = f"{cls.__name__}[{array_type.__name__}, '{dim_str}']"
+            # In python 3.8, e.g., typing.Union lacks `__name__`.
+            try:
+                type_str = array_type.__name__
+            except AttributeError:
+                type_str = repr(array_type)
+            name = f"{cls.__name__}[{type_str}, '{dim_str}']"
         elif _array_name_format == "array":
             name = array_type.__name__
         else:

setup.py

@@ -94,4 +94,5 @@ setuptools.setup(
     install_requires=install_requires,
     entry_points=entry_points,
     packages=[name],
+    include_package_data=True,
 )