From 066a5b058fe9d33537d5a87137d7df085ce9c6f6 Mon Sep 17 00:00:00 2001 From: Patrick Kidger <33688385+patrick-kidger@users.noreply.github.com> Date: Tue, 30 May 2023 12:34:10 -0700 Subject: [PATCH] Made modules private. --- jaxtyping/__init__.py | 18 +++++++++--------- jaxtyping/{array_types.py => _array_types.py} | 2 +- jaxtyping/{decorator.py => _decorator.py} | 0 jaxtyping/{import_hook.py => _import_hook.py} | 2 +- jaxtyping/{indirection.py => _indirection.py} | 0 .../{pytest_plugin.py => _pytest_plugin.py} | 2 +- jaxtyping/{pytree_type.py => _pytree_type.py} | 0 pyproject.toml | 4 ++-- 8 files changed, 14 insertions(+), 14 deletions(-) rename jaxtyping/{array_types.py => _array_types.py} (99%) rename jaxtyping/{decorator.py => _decorator.py} (100%) rename jaxtyping/{import_hook.py => _import_hook.py} (99%) rename jaxtyping/{indirection.py => _indirection.py} (100%) rename jaxtyping/{pytest_plugin.py => _pytest_plugin.py} (97%) rename jaxtyping/{pytree_type.py => _pytree_type.py} (100%) diff --git a/jaxtyping/__init__.py b/jaxtyping/__init__.py index 22ad5a8..1dbd5e0 100644 --- a/jaxtyping/__init__.py +++ b/jaxtyping/__init__.py @@ -30,14 +30,14 @@ else: del jax # First import some things as normal -from .array_types import ( +from ._array_types import ( AbstractArray as AbstractArray, AbstractDtype as AbstractDtype, get_array_name_format as get_array_name_format, set_array_name_format as set_array_name_format, ) -from .decorator import jaxtyped as jaxtyped -from .import_hook import install_import_hook as install_import_hook +from ._decorator import jaxtyped as jaxtyped +from ._import_hook import install_import_hook as install_import_hook # Now import Array and ArrayLike @@ -71,7 +71,7 @@ elif has_jax: if typing.TYPE_CHECKING: # Introduce an indirection so that we can `import X as X` to make it clear that # these are public. - from .indirection import ( + from ._indirection import ( BFloat16 as BFloat16, Bool as Bool, Complex as Complex, @@ -98,7 +98,7 @@ if typing.TYPE_CHECKING: UInt64 as UInt64, ) else: - from .array_types import ( + from ._array_types import ( BFloat16 as BFloat16, Bool as Bool, Complex as Complex, @@ -125,7 +125,7 @@ else: ) if has_jax: - from .array_types import Key as Key + from ._array_types import Key as Key # Now import PyTreeDef and PyTree @@ -155,16 +155,16 @@ if typing.TYPE_CHECKING: elif has_jax: from jax.tree_util import PyTreeDef as PyTreeDef - from .pytree_type import PyTree as PyTree # noqa: F401 + from ._pytree_type import PyTree as PyTree # noqa: F401 # Conveniences if typing.TYPE_CHECKING: from jax.random import PRNGKeyArray as PRNGKeyArray - from .indirection import Scalar as Scalar, ScalarLike as ScalarLike + from ._indirection import Scalar as Scalar, ScalarLike as ScalarLike elif has_jax: - from .array_types import PRNGKeyArray, Scalar, ScalarLike # noqa: F401 + from ._array_types import PRNGKeyArray, Scalar, ScalarLike # noqa: F401 del has_jax diff --git a/jaxtyping/array_types.py b/jaxtyping/_array_types.py similarity index 99% rename from jaxtyping/array_types.py rename to jaxtyping/_array_types.py index 489efbe..42bcbf9 100644 --- a/jaxtyping/array_types.py +++ b/jaxtyping/_array_types.py @@ -36,7 +36,7 @@ from typing import ( import numpy as np -from .decorator import storage +from ._decorator import storage try: diff --git a/jaxtyping/decorator.py b/jaxtyping/_decorator.py similarity index 100% rename from jaxtyping/decorator.py rename to jaxtyping/_decorator.py diff --git a/jaxtyping/import_hook.py b/jaxtyping/_import_hook.py similarity index 99% rename from jaxtyping/import_hook.py rename to jaxtyping/_import_hook.py index ab5497f..3c91a8f 100644 --- a/jaxtyping/import_hook.py +++ b/jaxtyping/_import_hook.py @@ -120,7 +120,7 @@ class _JaxtypingTransformer(ast.NodeVisitor): return node def visit_ClassDef(self, node: ast.ClassDef): - func = _dot_lookup("jaxtyping", "decorator", "_jaxtyped_typechecker") + func = _dot_lookup("jaxtyping", "_decorator", "_jaxtyped_typechecker") if self._typechecker is None: args = [ast.Constant(None)] else: diff --git a/jaxtyping/indirection.py b/jaxtyping/_indirection.py similarity index 100% rename from jaxtyping/indirection.py rename to jaxtyping/_indirection.py diff --git a/jaxtyping/pytest_plugin.py b/jaxtyping/_pytest_plugin.py similarity index 97% rename from jaxtyping/pytest_plugin.py rename to jaxtyping/_pytest_plugin.py index 1c458f1..3d9dcb6 100644 --- a/jaxtyping/pytest_plugin.py +++ b/jaxtyping/_pytest_plugin.py @@ -19,7 +19,7 @@ import sys -from .import_hook import install_import_hook +from ._import_hook import install_import_hook def pytest_addoption(parser): diff --git a/jaxtyping/pytree_type.py b/jaxtyping/_pytree_type.py similarity index 100% rename from jaxtyping/pytree_type.py rename to jaxtyping/_pytree_type.py diff --git a/pyproject.toml b/pyproject.toml index dcaf0a0..076f271 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,6 @@ [project] name = "jaxtyping" -version = "0.2.19" +version = "0.2.20" description = "Type annotations and runtime checking for shape and dtype of JAX arrays, and PyTrees." readme = "README.md" requires-python ="~=3.8" @@ -24,7 +24,7 @@ classifiers = [ ] urls = {repository = "https://github.com/google/jaxtyping" } dependencies = ["numpy>=1.20.0", "typeguard>=2.13.3", "typing_extensions>=3.7.4.1"] -entry-points = {pytest11 = {jaxtyping = "jaxtyping.pytest_plugin"}} +entry-points = {pytest11 = {jaxtyping = "jaxtyping._pytest_plugin"}} [build-system] requires = ["hatchling"]