From 7b3d9a2e9aef5e644b7eac75b0175cf598e23fe1 Mon Sep 17 00:00:00 2001
From: Brent Yi <yibrenth@gmail.com>
Date: Wed, 7 Dec 2022 17:21:33 -0800
Subject: [PATCH] Explicitly export names in `jaxtyping.*` (#49)

* Explicitly export names to make pyright happy

* Bump jax and jaxtyping versions

* Add note on `jaxtyping` names

* Remove __all__ from `array_types.py`

* Appease flake8

* Reduce import redundancy

* Fix capitalization
---
 .isort.cfg               |  1 +
 jaxtyping/__init__.py    | 67 +++++++++++++++++++++-------------------
 jaxtyping/array_types.py | 48 ++++++++++++++--------------
 3 files changed, 61 insertions(+), 55 deletions(-)

diff --git a/.isort.cfg b/.isort.cfg
index bdc8fe4..f9e6dd4 100644
--- a/.isort.cfg
+++ b/.isort.cfg
@@ -2,5 +2,6 @@
 force_alphabetical_sort_within_sections=true
 lines_after_imports=2
 profile=black
+combine_as_imports=True
 treat_comments_as_code=true
 extra_standard_library=typing_extensions
diff --git a/jaxtyping/__init__.py b/jaxtyping/__init__.py
index 7a00501..06578ab 100644
--- a/jaxtyping/__init__.py
+++ b/jaxtyping/__init__.py
@@ -30,8 +30,11 @@ else:
     del jax
 
 
+# Type checkers don't know which branch below will be executed.
 if typing.TYPE_CHECKING:
-    from jax.numpy import ndarray as Array
+    # For imports, we need to explicitly `import X as X` in order for Pyright to see
+    # them as public. See discussion at https://github.com/microsoft/pyright/issues/2277
+    from jax import Array as Array
 elif has_jax:
     if getattr(typing, "GENERATING_DOCUMENTATION", False):
 
@@ -40,39 +43,39 @@ elif has_jax:
 
         Array.__module__ = "builtins"
     else:
-        from jax.numpy import ndarray as Array
+        from jax import Array as Array
 
 from .array_types import (
-    AbstractArray,
-    AbstractDtype,
-    BFloat16,
-    Bool,
-    Complex,
-    Complex64,
-    Complex128,
-    Float,
-    Float16,
-    Float32,
-    Float64,
-    get_array_name_format,
-    Inexact,
-    Int,
-    Int8,
-    Int16,
-    Int32,
-    Int64,
-    Integer,
-    Num,
-    set_array_name_format,
-    Shaped,
-    UInt,
-    UInt8,
-    UInt16,
-    UInt32,
-    UInt64,
+    AbstractArray as AbstractArray,
+    AbstractDtype as AbstractDtype,
+    BFloat16 as BFloat16,
+    Bool as Bool,
+    Complex as Complex,
+    Complex64 as Complex64,
+    Complex128 as Complex128,
+    Float as Float,
+    Float16 as Float16,
+    Float32 as Float32,
+    Float64 as Float64,
+    get_array_name_format as get_array_name_format,
+    Inexact as Inexact,
+    Int as Int,
+    Int8 as Int8,
+    Int16 as Int16,
+    Int32 as Int32,
+    Int64 as Int64,
+    Integer as Integer,
+    Num as Num,
+    set_array_name_format as set_array_name_format,
+    Shaped as Shaped,
+    UInt as UInt,
+    UInt8 as UInt8,
+    UInt16 as UInt16,
+    UInt32 as UInt32,
+    UInt64 as UInt64,
 )
-from .decorator import jaxtyped
-from .import_hook import install_import_hook
+from .decorator import jaxtyped as jaxtyped
+from .import_hook import install_import_hook as install_import_hook
 
 
 if typing.TYPE_CHECKING:
@@ -86,4 +89,4 @@ elif has_jax:
 
 del has_jax
 
-__version__ = "0.2.8"
+__version__ = "0.2.9"
diff --git a/jaxtyping/array_types.py b/jaxtyping/array_types.py
index 32790c4..f218269 100644
--- a/jaxtyping/array_types.py
+++ b/jaxtyping/array_types.py
@@ -433,29 +433,31 @@ if TYPE_CHECKING:
     # Note that `from typing_extensions import Annotated; ... = Annotated`
     # does not work with static type checkers. `Annotated` is a typeform rather
     # than a type, meaning it cannot be assigned.
-    from typing_extensions import Annotated as BFloat16
-    from typing_extensions import Annotated as Bool
-    from typing_extensions import Annotated as Complex
-    from typing_extensions import Annotated as Complex64
-    from typing_extensions import Annotated as Complex128
-    from typing_extensions import Annotated as Float
-    from typing_extensions import Annotated as Float16
-    from typing_extensions import Annotated as Float32
-    from typing_extensions import Annotated as Float64
-    from typing_extensions import Annotated as Inexact
-    from typing_extensions import Annotated as Int
-    from typing_extensions import Annotated as Int8
-    from typing_extensions import Annotated as Int16
-    from typing_extensions import Annotated as Int32
-    from typing_extensions import Annotated as Int64
-    from typing_extensions import Annotated as Integer
-    from typing_extensions import Annotated as Num
-    from typing_extensions import Annotated as Shaped
-    from typing_extensions import Annotated as UInt
-    from typing_extensions import Annotated as UInt8
-    from typing_extensions import Annotated as UInt16
-    from typing_extensions import Annotated as UInt32
-    from typing_extensions import Annotated as UInt64
+    from typing_extensions import (
+        Annotated as BFloat16,
+        Annotated as Bool,
+        Annotated as Complex,
+        Annotated as Complex64,
+        Annotated as Complex128,
+        Annotated as Float,
+        Annotated as Float16,
+        Annotated as Float32,
+        Annotated as Float64,
+        Annotated as Inexact,
+        Annotated as Int,
+        Annotated as Int8,
+        Annotated as Int16,
+        Annotated as Int32,
+        Annotated as Int64,
+        Annotated as Integer,
+        Annotated as Num,
+        Annotated as Shaped,
+        Annotated as UInt,
+        Annotated as UInt8,
+        Annotated as UInt16,
+        Annotated as UInt32,
+        Annotated as UInt64,
+    )
 else:
     _bool = "bool_"
     _uint8 = "uint8"