diff --git a/.isort.cfg b/.isort.cfg index 64737cc..bdc8fe4 100644 --- a/.isort.cfg +++ b/.isort.cfg @@ -3,3 +3,4 @@ force_alphabetical_sort_within_sections=true lines_after_imports=2 profile=black treat_comments_as_code=true +extra_standard_library=typing_extensions diff --git a/jaxtyping/__init__.py b/jaxtyping/__init__.py index 8547022..bf4a5ad 100644 --- a/jaxtyping/__init__.py +++ b/jaxtyping/__init__.py @@ -51,4 +51,4 @@ from .import_hook import install_import_hook from .pytree_type import PyTree -__version__ = "0.0.1" +__version__ = "0.0.2" diff --git a/jaxtyping/array_types.py b/jaxtyping/array_types.py index 667a090..348ee06 100644 --- a/jaxtyping/array_types.py +++ b/jaxtyping/array_types.py @@ -18,7 +18,8 @@ # CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. import functools as ft -from typing import Any, Dict, List, Literal, NoReturn, Optional, Tuple, Union +from typing import Any, Dict, List, NoReturn, Optional, Tuple, Union +from typing_extensions import Literal import jax.numpy as jnp diff --git a/jaxtyping/pytree_type.py b/jaxtyping/pytree_type.py index 7a7d894..c19a0bc 100644 --- a/jaxtyping/pytree_type.py +++ b/jaxtyping/pytree_type.py @@ -18,11 +18,27 @@ # CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. import functools as ft +from typing import Generic, TypeVar import jax import typeguard +_T = TypeVar("_T") + + +class _FakePyTree(Generic[_T]): + pass + + +_FakePyTree.__name__ = "PyTree" +_FakePyTree.__qualname__ = "PyTree" +# Can't do type("PyTree", (Generic[_T],), {}) because dynamic subclassing of typeforms +# isn't allowed. +# Can't do types.new_class("PyTree", (Generic[_T],), {}) because that has __module__ +# "types", e.g. we get types.PyTree[int]. + + class _MetaPyTree(type): def __call__(self, *args, **kwargs): raise RuntimeError("PyTree cannot be instantiated") @@ -32,7 +48,8 @@ class _MetaPyTree(type): @ft.lru_cache(maxsize=None) def __getitem__(cls, item): - return _MetaSubscriptPyTree(f"PyTree[{item.__name__}]", (), {"leaftype": item}) + name = str(_FakePyTree[item]) + return _MetaSubscriptPyTree(name, (), {"leaftype": item}) class _MetaSubscriptPyTree(type): @@ -63,3 +80,6 @@ class _MetaSubscriptPyTree(type): PyTree = _MetaPyTree("PyTree", (), {}) +# Can't do `class PyTree(Generic[_T]): ...` because we need to override the +# instancecheck for PyTree[foo], but we subclassing +# `type(Generic[int])`, i.e. `typing._GenericAlias` is disallowed. diff --git a/setup.py b/setup.py index 9776fcd..1a97619 100644 --- a/setup.py +++ b/setup.py @@ -63,7 +63,7 @@ python_requires = "~=3.7" # We use typeguard internally (in a fairly minimal way), but it's not required that # end users make the same choice. -install_requires = ["jax>=0.3.4", "typeguard>=2.13.3"] +install_requires = ["jax>=0.3.4", "typeguard>=2.13.3", "typing_extensions>=4.2.0"] entry_points = dict(pytest11=["jaxtyping = jaxtyping.pytest_plugin"])