Hopefully fixed PyTree raising spurious errors. Bit mysterious that this worked before, really. I've tested this fix as best I can against the various static type checkers, but these are weird and varied enough that this might not be a perfect fix. If you see this and have issues, let me know. (#54)

This commit is contained in:
Patrick Kidger
2022-12-30 19:00:26 +00:00
committed by GitHub
parent 7dba3516c2
commit 59e8fb0d18
+19 -6
View File
@@ -79,14 +79,27 @@ from .import_hook import install_import_hook as install_import_hook
if typing.TYPE_CHECKING:
_T = typing.TypeVar("_T")
class PyTree(typing_extensions.Protocol[_T]):
pass
# Set up to deliberately confuse a static type checker.
PyTree = getattr(typing, "foo" + "bar")
# What's going on with this madness?
#
# At static-type-checking-time, we want `PyTree` to be a type for which both
# `PyTree` and `PyTree[Foo]` are equivalent to `Any`.
# (The intention is that `PyTree` be a runtime-only type; there's no real way to
# do more with static type checkers.)
#
# Unfortunately, this isn't possible: `Any` isn't subscriptable. And there's no
# equivalent way we can fake this using typing annotations. (In some sense the
# closest thing would be a `Protocol[T]` with no methods, but that's actually the
# opposite of what we want: that ends up allowing nothing at all.)
#
# The good news for us is that static type checkers have an internal escape hatch.
# If they can't figure out what a type is, then they just give up and allow
# anything. (I believe this is sometimes called `Unknown`.) Thus, this odd-looking
# annotation, which static type checkers aren't smart enough to resolve.
elif has_jax:
from .pytree_type import PyTree
del has_jax
__version__ = "0.2.10"
__version__ = "0.2.11"