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.

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.5"
+__version__ = "0.2.6"

jaxtyping/array_types.py

@@ -384,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: