From dcd73e343124d1ed9d2bbd356a50fef6a60a45c9 Mon Sep 17 00:00:00 2001
From: ebrevdo <ebrevdo@users.noreply.github.com>
Date: Thu, 22 Sep 2022 13:25:29 -0700
Subject: [PATCH] Add support for e.g. jaxtyping.Float[Union[...], ...] in
 py3.8 (#31)

* 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/array_types.py | 7 ++++++-
 1 file changed, 6 insertions(+), 1 deletion(-)

diff --git a/jaxtyping/array_types.py b/jaxtyping/array_types.py
index 2e18b1d..c96a5f0 100644
--- a/jaxtyping/array_types.py
+++ b/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: