diff --git a/.clang-format b/.clang-format index 8f6cc725..bf0f6ced 100644 --- a/.clang-format +++ b/.clang-format @@ -13,8 +13,10 @@ AlignEscapedNewlines: Right AllowAllArgumentsOnNextLine: false AllowAllParametersOfDeclarationOnNextLine: false AllowShortIfStatementsOnASingleLine: Never -BinPackArguments: false -BinPackParameters: OnePerLine +PackArguments: + BinPack: OnePerLine +PackParameters: + BinPack: OnePerLine BreakBeforeTernaryOperators: true CommentPragmas: 'NOLINT(NEXTLINE|BEGIN|END)?\[.*\]' DerivePointerAlignment: false diff --git a/.github/workflows/build.yml b/.github/workflows/build.yml index 79b7a55c..bc834082 100644 --- a/.github/workflows/build.yml +++ b/.github/workflows/build.yml @@ -82,7 +82,7 @@ jobs: - name: Set up Python uses: actions/setup-python@v7 with: - python-version: "3.9 - 3.15" + python-version: "3.10 - 3.15" update-environment: true - name: Upgrade pip @@ -151,7 +151,6 @@ jobs: - { runner: windows-latest, platform: windows, archs: "auto32" } - { runner: windows-11-arm, platform: windows, archs: "ARM64" } python-version: - - "3.9" - "3.10" - "3.11" - "3.12" @@ -164,8 +163,6 @@ jobs: - "pypy-3.11" exclude: # Exclude unsupported Python versions - - python-version: "3.9" - target: { archs: "ARM64" } - python-version: "3.10" target: { archs: "ARM64" } - python-version: "pypy-3.11" @@ -429,7 +426,7 @@ jobs: if: startsWith(github.ref, 'refs/tags/') uses: actions/setup-python@v7 with: - python-version: "3.9 - 3.15" + python-version: "3.10 - 3.15" update-environment: true - name: Upgrade pip diff --git a/.github/workflows/tests-with-pydebug.yml b/.github/workflows/tests-with-pydebug.yml index 926b67ba..8660d131 100644 --- a/.github/workflows/tests-with-pydebug.yml +++ b/.github/workflows/tests-with-pydebug.yml @@ -70,7 +70,6 @@ jobs: matrix: runner: [ubuntu-latest, macos-latest, windows-latest] python-version: - - "3.9" - "3.10" - "3.11" - "3.12" @@ -79,8 +78,6 @@ jobs: - "3.15" python-abiflags: ["d", "td"] exclude: - - python-version: "3.9" - python-abiflags: "td" - python-version: "3.10" python-abiflags: "td" - python-version: "3.11" diff --git a/.github/workflows/tests.yml b/.github/workflows/tests.yml index c9a979c1..9ee88c52 100644 --- a/.github/workflows/tests.yml +++ b/.github/workflows/tests.yml @@ -71,7 +71,6 @@ jobs: matrix: runner: [ubuntu-latest, macos-latest, windows-latest] python-version: - - "3.9" - "3.10" - "3.11" - "3.12" diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 0d0545ea..8e40fb84 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -30,7 +30,7 @@ repos: args: [--ignore-case] files: ^docs/source/spelling_wordlist\.txt$ - repo: https://github.com/pre-commit/mirrors-clang-format - rev: v22.1.8 + rev: v23.1.0 hooks: - id: clang-format - repo: https://github.com/cpplint/cpplint @@ -38,7 +38,7 @@ repos: hooks: - id: cpplint - repo: https://github.com/astral-sh/ruff-pre-commit - rev: v0.16.3 + rev: v0.16.5 hooks: - id: ruff-check args: [--fix, --exit-non-zero-on-fix] diff --git a/CHANGELOG.md b/CHANGELOG.md index 8262b1d2..b3ffabcb 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -25,7 +25,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ### Removed -- +- Drop Python 3.9 support by [@XuehaiPan](https://github.com/XuehaiPan) in [#272](https://github.com/metaopt/optree/pull/272). ------ diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index 439104cc..2f08a271 100644 --- a/CONTRIBUTING.md +++ b/CONTRIBUTING.md @@ -88,13 +88,13 @@ python3 -m cibuildwheel --platform=linux --output-dir=wheelhouse --config-file=p It will build wheel binaries for all supported CPython versions. The outputs will be placed in the `wheelhouse` directory. To build a wheel for a specific CPython version, you can use the [`CIBW_BUILD`](https://cibuildwheel.readthedocs.io/en/stable/options/#build-skip) environment variable. -For example, the following command will build a wheel for Python 3.9: +For example, the following command will build a wheel for Python 3.10: ```bash -CIBW_BUILD="cp39*manylinux*" python3 -m cibuildwheel --platform=linux --output-dir=wheelhouse --config-file=pyproject.toml +CIBW_BUILD="cp310*manylinux*" python3 -m cibuildwheel --platform=linux --output-dir=wheelhouse --config-file=pyproject.toml ``` -You can change `cp39*` to `cp312*` to build for Python 3.12. See for more options. +You can change `cp310*` to `cp312*` to build for Python 3.12. See for more options. ## Documentation diff --git a/README.md b/README.md index acac4baf..b32b7862 100644 --- a/README.md +++ b/README.md @@ -2,7 +2,7 @@ # OpTree -![Python 3.9+](https://img.shields.io/badge/Python-3.9%2B-brightgreen) +![Python 3.10+](https://img.shields.io/badge/Python-3.10%2B-brightgreen) [![PyPI](https://img.shields.io/pypi/v/optree?logo=pypi)](https://pypi.org/project/optree) ![GitHub Workflow Status](https://img.shields.io/github/actions/workflow/status/metaopt/optree/build.yml?label=build&logo=github) ![GitHub Workflow Status](https://img.shields.io/github/actions/workflow/status/metaopt/optree/tests.yml?label=tests&logo=github) @@ -72,7 +72,7 @@ export pybind11_DIR="/path/to/custom/pybind11" pip3 install . ``` -Compiling from source requires Python 3.9+, a C++ compiler (`g++` / `clang++` / `icpx` / `cl.exe`) that supports C++20, and a `cmake` installation. +Compiling from source requires Python 3.10+, a C++ compiler (`g++` / `clang++` / `icpx` / `cl.exe`) that supports C++20, and a `cmake` installation. -------------------------------------------------------------------------------- @@ -240,7 +240,7 @@ optree.register_pytree_node( list(vars(ct).keys()), ), unflatten_func=lambda keys, values: ( # unflatten: (metadata, children) -> MyContainer - MyContainer(**dict(zip(keys, values))) + MyContainer(**dict(zip(keys, values, strict=True))) ), path_entry_type=MyContainerEntry, namespace='mycontainer', @@ -305,7 +305,7 @@ class MyDict(UserDict): @classmethod def __tree_unflatten__(cls, metadata, children): - return cls(zip(metadata, children)) + return cls(zip(metadata, children, strict=True)) ``` ```python @@ -389,7 +389,7 @@ There are several key attributes of the pytree type registry: @classmethod def __tree_unflatten__(cls, metadata, children): - return cls(zip(metadata, children)) + return cls(zip(metadata, children, strict=True)) # Subclasses will be automatically registered in namespace 'mydict' diff --git a/docs/source/conf.py b/docs/source/conf.py index 7a5adfa0..b0c331c4 100644 --- a/docs/source/conf.py +++ b/docs/source/conf.py @@ -223,7 +223,7 @@ def get_version() -> str: def get_pytree_typing_instance(annotation): # noqa: C901 if not ( - isinstance(annotation, type(typing.Union[int, str])) + isinstance(annotation, type(typing.Union[int, str])) # noqa: UP007 and typing.get_origin(annotation) is typing.Union ): return None @@ -281,7 +281,7 @@ def matches_pytree_typing_alias( return all( matches_pytree_typing_alias(arg, pat, recursive_pattern, recursive_ref_names) - for arg, pat in zip(annotation_args, pattern_args) + for arg, pat in zip(annotation_args, pattern_args, strict=True) ) for pytree_alias, pytree_instance in tuple(PyTree.__instances__.items()): @@ -321,7 +321,7 @@ def format_pytree_typing_union(annotation_args, config=None): pytree_instance = None next_index = index for end in range(len(annotation_args), index + min_pytree_union_width - 1, -1): - pytree_instance = get_pytree_typing_instance(typing.Union[annotation_args[index:end]]) + pytree_instance = get_pytree_typing_instance(typing.Union[annotation_args[index:end]]) # noqa: UP007 if pytree_instance is not None: next_index = end break diff --git a/include/optree/pymacros.h b/include/optree/pymacros.h index 59bd705f..261462a8 100644 --- a/include/optree/pymacros.h +++ b/include/optree/pymacros.h @@ -23,8 +23,8 @@ limitations under the License. #include -#if !(defined(PY_VERSION_HEX) && PY_VERSION_HEX >= 0x03090000) // Python 3.9 -# error "Python 3.9 or newer is required." +#if !(defined(PY_VERSION_HEX) && PY_VERSION_HEX >= 0x030A0000) // Python 3.10 +# error "Python 3.10 or newer is required." #endif #if !(defined(PYBIND11_VERSION_HEX) && PYBIND11_VERSION_HEX >= 0x020C00F0) // pybind11 2.12.0 diff --git a/include/optree/treespec.h b/include/optree/treespec.h index 35505620..9b4b5863 100644 --- a/include/optree/treespec.h +++ b/include/optree/treespec.h @@ -46,11 +46,7 @@ using ssize_t = py::ssize_t; #if defined(MS_WINDOWS) && (defined(Py_DEBUG) || defined(Py_GIL_DISABLED)) // A debug or free-threading build on Windows combines large frames with a 1MB default stack, and // the walkers still need room to unwind the `RecursionError` thrown at the guard. -# if PY_VERSION_HEX < 0x030A0000 // Python 3.10 -constexpr ssize_t MAX_RECURSION_DEPTH = 100; -# else constexpr ssize_t MAX_RECURSION_DEPTH = 250; -# endif #elif defined(Py_DEBUG) || defined(PYPY_VERSION) || defined(MS_WINDOWS) || \ (defined(__wasm__) || defined(__wasm32__) || defined(__wasm64__) || defined(__wasi__) || \ defined(__EMSCRIPTEN__)) diff --git a/optree/accessors.py b/optree/accessors.py index 6fcd2399..141bea7b 100644 --- a/optree/accessors.py +++ b/optree/accessors.py @@ -17,7 +17,6 @@ from __future__ import annotations import dataclasses -import sys from collections.abc import Iterable, Mapping, Sequence from typing import TYPE_CHECKING, Any, ClassVar, Generic, Literal, TypeVar, overload from typing_extensions import Self # Python 3.11+ @@ -47,10 +46,7 @@ ] -SLOTS = {'slots': True} if sys.version_info >= (3, 10) else {} # Python 3.10+ - - -@dataclasses.dataclass(init=True, repr=False, eq=False, frozen=True, **SLOTS) +@dataclasses.dataclass(init=True, repr=False, eq=False, frozen=True, slots=True) class PyTreeEntry: """Base class for path entries.""" @@ -122,9 +118,6 @@ def codify(self, /, node: str = '') -> str: return f'{node}[]' # should be overridden -del SLOTS - - _T = TypeVar('_T') _T_co = TypeVar('_T_co', covariant=True) _KT_co = TypeVar('_KT_co', covariant=True) @@ -134,7 +127,7 @@ def codify(self, /, node: str = '') -> str: class AutoEntry(PyTreeEntry): """A generic path entry class that determines the entry type on creation automatically.""" - __slots__: ClassVar[tuple[()]] = () + __slots__: ClassVar[tuple[()]] = () # type: ignore[misc] def __new__( # type: ignore[misc] cls, @@ -184,7 +177,7 @@ def __new__( # type: ignore[misc] class GetItemEntry(PyTreeEntry): """A generic path entry class for nodes that access their children by :meth:`__getitem__`.""" - __slots__: ClassVar[tuple[()]] = () + __slots__: ClassVar[tuple[()]] = () # type: ignore[misc] def __call__(self, obj: Any, /) -> Any: """Get the child object.""" @@ -198,7 +191,7 @@ def codify(self, /, node: str = '') -> str: class GetAttrEntry(PyTreeEntry): """A generic path entry class for nodes that access their children by :meth:`__getattr__`.""" - __slots__: ClassVar[tuple[()]] = () + __slots__: ClassVar[tuple[()]] = () # type: ignore[misc] entry: str @@ -223,13 +216,13 @@ def codify(self, /, node: str = '') -> str: class FlattenedEntry(PyTreeEntry): # pylint: disable=too-few-public-methods """A fallback path entry class for flattened objects.""" - __slots__: ClassVar[tuple[()]] = () + __slots__: ClassVar[tuple[()]] = () # type: ignore[misc] class SequenceEntry(GetItemEntry, Generic[_T_co]): """A path entry class for sequences.""" - __slots__: ClassVar[tuple[()]] = () + __slots__: ClassVar[tuple[()]] = () # type: ignore[misc] entry: int type: builtins.type[Sequence[_T_co]] @@ -251,7 +244,7 @@ def __repr__(self, /) -> str: class MappingEntry(GetItemEntry, Generic[_KT_co, _VT_co]): """A path entry class for mappings.""" - __slots__: ClassVar[tuple[()]] = () + __slots__: ClassVar[tuple[()]] = () # type: ignore[misc] entry: _KT_co type: builtins.type[Mapping[_KT_co, _VT_co]] @@ -273,7 +266,7 @@ def __repr__(self, /) -> str: class NamedTupleEntry(SequenceEntry[_T]): """A path entry class for namedtuple objects.""" - __slots__: ClassVar[tuple[()]] = () + __slots__: ClassVar[tuple[()]] = () # type: ignore[misc] entry: int type: builtins.type[NamedTuple[_T]] # type: ignore[type-arg] @@ -303,7 +296,7 @@ def codify(self, /, node: str = '') -> str: class StructSequenceEntry(SequenceEntry[_T]): """A path entry class for PyStructSequence objects.""" - __slots__: ClassVar[tuple[()]] = () + __slots__: ClassVar[tuple[()]] = () # type: ignore[misc] entry: int type: builtins.type[StructSequence[_T]] @@ -347,7 +340,7 @@ def codify(self, /, node: str = '') -> str: class DataclassEntry(GetAttrEntry): """A path entry class for dataclasses.""" - __slots__: ClassVar[tuple[()]] = () + __slots__: ClassVar[tuple[()]] = () # type: ignore[misc] entry: str | int # type: ignore[assignment] diff --git a/optree/dataclasses.py b/optree/dataclasses.py index e41b653f..244a0ae1 100644 --- a/optree/dataclasses.py +++ b/optree/dataclasses.py @@ -69,14 +69,14 @@ import warnings from dataclasses import * # noqa: F401,F403,RUF100 # pylint: disable=wildcard-import,unused-wildcard-import from types import MappingProxyType -from typing import TYPE_CHECKING, Any, Callable, Literal, Protocol, TypeVar, overload +from typing import TYPE_CHECKING, Any, Literal, Protocol, TypeVar, overload from typing_extensions import dataclass_transform # Python 3.11+ from optree.accessors import DataclassEntry if TYPE_CHECKING: - from collections.abc import Iterable + from collections.abc import Callable, Iterable __all__ = [ @@ -106,7 +106,7 @@ def field( hash: bool | None = None, # pylint: disable=redefined-builtin compare: bool = True, metadata: dict[Any, Any] | None = None, - kw_only: bool | Literal[dataclasses.MISSING] = dataclasses.MISSING, # type: ignore[valid-type] # Python 3.10+ + kw_only: bool | Literal[dataclasses.MISSING] = dataclasses.MISSING, # type: ignore[valid-type] doc: str | None = None, # Python 3.14+ pytree_node: bool | None = None, ) -> _T: ... @@ -121,7 +121,7 @@ def field( hash: bool | None = None, # pylint: disable=redefined-builtin compare: bool = True, metadata: dict[Any, Any] | None = None, - kw_only: bool | Literal[dataclasses.MISSING] = dataclasses.MISSING, # type: ignore[valid-type] # Python 3.10+ + kw_only: bool | Literal[dataclasses.MISSING] = dataclasses.MISSING, # type: ignore[valid-type] doc: str | None = None, # Python 3.14+ pytree_node: bool | None = None, ) -> _T: ... @@ -135,7 +135,7 @@ def field( hash: bool | None = None, # pylint: disable=redefined-builtin compare: bool = True, metadata: dict[Any, Any] | None = None, - kw_only: bool | Literal[dataclasses.MISSING] = dataclasses.MISSING, # type: ignore[valid-type] # Python 3.10+ + kw_only: bool | Literal[dataclasses.MISSING] = dataclasses.MISSING, # type: ignore[valid-type] doc: str | None = None, # Python 3.14+ pytree_node: bool | None = None, ) -> Any: ... @@ -150,7 +150,7 @@ def field( # noqa: D417 # pylint: disable=function-redefined hash: bool | None = None, # pylint: disable=redefined-builtin compare: bool = True, metadata: dict[Any, Any] | None = None, - kw_only: bool | Literal[dataclasses.MISSING] = dataclasses.MISSING, # type: ignore[valid-type] # Python 3.10+ + kw_only: bool | Literal[dataclasses.MISSING] = dataclasses.MISSING, # type: ignore[valid-type] doc: str | None = None, # Python 3.14+ pytree_node: bool | None = None, ) -> Any: @@ -192,13 +192,9 @@ def field( # noqa: D417 # pylint: disable=function-redefined 'hash': hash, 'compare': compare, 'metadata': metadata, + 'kw_only': kw_only, } - if sys.version_info >= (3, 10): # pragma: >=3.10 cover - kwargs['kw_only'] = kw_only - elif kw_only is not dataclasses.MISSING: # pragma: <3.10 cover - raise TypeError("field() got an unexpected keyword argument 'kw_only'") - if sys.version_info >= (3, 14): # pragma: >=3.14 cover kwargs['doc'] = doc elif doc is not None: # pragma: <3.14 cover @@ -222,9 +218,9 @@ def dataclass( order: bool = False, unsafe_hash: bool = False, frozen: bool = False, - match_args: bool = True, # Python 3.10+ - kw_only: bool = False, # Python 3.10+ - slots: bool = False, # Python 3.10+ + match_args: bool = True, + kw_only: bool = False, + slots: bool = False, weakref_slot: bool = False, # Python 3.11+ namespace: str, ) -> Callable[[_TypeT], _TypeT]: ... @@ -241,16 +237,16 @@ def dataclass( order: bool = False, unsafe_hash: bool = False, frozen: bool = False, - match_args: bool = True, # Python 3.10+ - kw_only: bool = False, # Python 3.10+ - slots: bool = False, # Python 3.10+ + match_args: bool = True, + kw_only: bool = False, + slots: bool = False, weakref_slot: bool = False, # Python 3.11+ namespace: str, ) -> _TypeT: ... @dataclass_transform(field_specifiers=(field,)) -def dataclass( # noqa: C901,D417 # pylint: disable=function-redefined +def dataclass( # noqa: D417 # pylint: disable=function-redefined cls: _TypeT | None = None, /, *, @@ -260,9 +256,9 @@ def dataclass( # noqa: C901,D417 # pylint: disable=function-redefined order: bool = False, unsafe_hash: bool = False, frozen: bool = False, - match_args: bool = True, # Python 3.10+ - kw_only: bool = False, # Python 3.10+ - slots: bool = False, # Python 3.10+ + match_args: bool = True, + kw_only: bool = False, + slots: bool = False, weakref_slot: bool = False, # Python 3.11+ namespace: str, ) -> _TypeT | Callable[[_TypeT], _TypeT]: @@ -289,19 +285,11 @@ def dataclass( # noqa: C901,D417 # pylint: disable=function-redefined 'order': order, 'unsafe_hash': unsafe_hash, 'frozen': frozen, + 'match_args': match_args, + 'kw_only': kw_only, + 'slots': slots, } - if sys.version_info >= (3, 10): # pragma: >=3.10 cover - kwargs['match_args'] = match_args - kwargs['kw_only'] = kw_only - kwargs['slots'] = slots - elif match_args is not True: # pragma: <3.10 cover - raise TypeError("dataclass() got an unexpected keyword argument 'match_args'") - elif kw_only is not False: # pragma: <3.10 cover - raise TypeError("dataclass() got an unexpected keyword argument 'kw_only'") - elif slots is not False: # pragma: <3.10 cover - raise TypeError("dataclass() got an unexpected keyword argument 'slots'") - if sys.version_info >= (3, 11): # pragma: >=3.11 cover kwargs['weakref_slot'] = weakref_slot elif weakref_slot is not False: # pragma: <3.11 cover @@ -363,9 +351,9 @@ def make_dataclass( # type: ignore[no-redef] # noqa: C901,D417 order: bool = False, unsafe_hash: bool = False, frozen: bool = False, - match_args: bool = True, # Python 3.10+ - kw_only: bool = False, # Python 3.10+ - slots: bool = False, # Python 3.10+ + match_args: bool = True, + kw_only: bool = False, + slots: bool = False, weakref_slot: bool = False, # Python 3.11+ module: str | None = None, # Python 3.12+ decorator: _DataclassDecorator[_TypeT] = dataclasses.dataclass, # type: ignore[assignment] # Python 3.14+ @@ -417,23 +405,15 @@ def make_dataclass( # type: ignore[no-redef] # noqa: C901,D417 'order': order, 'unsafe_hash': unsafe_hash, 'frozen': frozen, + 'match_args': match_args, + 'kw_only': kw_only, + 'slots': slots, } make_dataclass_kwargs = { 'bases': bases, 'namespace': ns, } - if sys.version_info >= (3, 10): # pragma: >=3.10 cover - dataclass_kwargs['match_args'] = match_args - dataclass_kwargs['kw_only'] = kw_only - dataclass_kwargs['slots'] = slots - elif match_args is not True: # pragma: <3.10 cover - raise TypeError("make_dataclass() got an unexpected keyword argument 'match_args'") - elif kw_only is not False: # pragma: <3.10 cover - raise TypeError("make_dataclass() got an unexpected keyword argument 'kw_only'") - elif slots is not False: # pragma: <3.10 cover - raise TypeError("make_dataclass() got an unexpected keyword argument 'slots'") - if sys.version_info >= (3, 11): # pragma: >=3.11 cover dataclass_kwargs['weakref_slot'] = weakref_slot elif weakref_slot is not False: # pragma: <3.11 cover @@ -626,7 +606,7 @@ def flatten_func( # pylint: disable-next=line-too-long def unflatten_func(metadata: tuple[tuple[str, Any], ...], children: tuple[_U, ...], /) -> _T: # type: ignore[type-var] - kwargs = dict(zip(children_field_names, children)) + kwargs = dict(zip(children_field_names, children, strict=True)) kwargs.update(metadata) return cls(**kwargs) # type: ignore[return-value] diff --git a/optree/functools.py b/optree/functools.py index 21563d28..edbefaa8 100644 --- a/optree/functools.py +++ b/optree/functools.py @@ -18,7 +18,7 @@ import contextlib import functools -from typing import TYPE_CHECKING, Any, Callable, ClassVar +from typing import TYPE_CHECKING, Any, ClassVar from typing_extensions import Self # Python 3.11+ from optree import registry @@ -28,6 +28,8 @@ if TYPE_CHECKING: + from collections.abc import Callable + from optree.accessors import PyTreeEntry @@ -141,10 +143,10 @@ class partial( # noqa: N801 # pylint: disable=invalid-name,too-few-public-metho def __new__(cls, func: Callable[..., Any], /, *args: T, **keywords: T) -> Self: """Create a new :class:`partial` instance.""" - # In Python 3.10+, if func is itself a functools.partial instance, functools.partial.__new__ - # would merge the arguments of this partial instance with the arguments of the func. We box - # func in a class that does not (yet) have a `func` attribute to defeat this optimization, - # since we care exactly which arguments are considered part of the pytree. + # If func is itself a functools.partial instance, functools.partial.__new__ merges its + # arguments into this one, and since Python 3.10 it does so for a subclass such as this one + # too. We box func in a class that does not (yet) have a `func` attribute to defeat this + # optimization, since we care exactly which arguments are considered part of the pytree. if isinstance(func, functools.partial): original_func = func func = _HashablePartialShim(original_func) diff --git a/optree/integrations/attrs.py b/optree/integrations/attrs.py index 98ad6b02..c881da3e 100644 --- a/optree/integrations/attrs.py +++ b/optree/integrations/attrs.py @@ -56,7 +56,7 @@ import inspect import warnings from types import MappingProxyType -from typing import TYPE_CHECKING, Any, Callable, TypeVar, overload +from typing import TYPE_CHECKING, Any, TypeVar, overload import attrs from attrs import ( @@ -84,6 +84,7 @@ if TYPE_CHECKING: + from collections.abc import Callable from typing import ClassVar @@ -129,7 +130,7 @@ class AttrsEntry(GetAttrEntry): """A path entry class for attrs classes.""" - __slots__: ClassVar[tuple[()]] = () + __slots__: ClassVar[tuple[()]] = () # type: ignore[misc] entry: str | int # type: ignore[assignment] @@ -491,7 +492,7 @@ def flatten_func( # pylint: disable-next=line-too-long def unflatten_func(metadata: tuple[tuple[str, Any], ...], children: tuple[_U, ...], /) -> _T: # type: ignore[type-var] - kwargs = dict(zip(children_aliases, children)) + kwargs = dict(zip(children_aliases, children, strict=True)) kwargs.update(metadata) return cls(**kwargs) # type: ignore[return-value] diff --git a/optree/integrations/jax.py b/optree/integrations/jax.py index 9b576149..84bf0a67 100644 --- a/optree/integrations/jax.py +++ b/optree/integrations/jax.py @@ -41,8 +41,7 @@ import warnings from operator import itemgetter from types import FunctionType -from typing import Any, Callable -from typing_extensions import TypeAlias # Python 3.10+ +from typing import TYPE_CHECKING, Any, TypeAlias import jax.numpy as jnp from jax import Array, lax @@ -54,6 +53,10 @@ from optree.utils import safe_zip, total_order_sorted +if TYPE_CHECKING: + from collections.abc import Callable + + __all__ = ['ArrayLikeTree', 'ArrayTree', 'tree_ravel'] diff --git a/optree/integrations/numpy.py b/optree/integrations/numpy.py index 4a5d0a9d..bb6431ce 100644 --- a/optree/integrations/numpy.py +++ b/optree/integrations/numpy.py @@ -22,8 +22,7 @@ import functools import itertools import warnings -from typing import Any, Callable -from typing_extensions import TypeAlias # Python 3.10+ +from typing import TYPE_CHECKING, Any, TypeAlias import numpy as np from numpy.typing import ArrayLike @@ -33,6 +32,10 @@ from optree.utils import safe_zip +if TYPE_CHECKING: + from collections.abc import Callable + + __all__ = ['ArrayLikeTree', 'ArrayTree', 'tree_ravel'] diff --git a/optree/integrations/torch.py b/optree/integrations/torch.py index cf8a1a03..dccbec84 100644 --- a/optree/integrations/torch.py +++ b/optree/integrations/torch.py @@ -21,8 +21,7 @@ import functools import warnings -from typing import Any, Callable -from typing_extensions import TypeAlias # Python 3.10+ +from typing import TYPE_CHECKING, Any, TypeAlias import torch @@ -31,6 +30,10 @@ from optree.utils import safe_zip +if TYPE_CHECKING: + from collections.abc import Callable + + __all__ = ['TensorTree', 'tree_ravel'] diff --git a/optree/ops.py b/optree/ops.py index 49bcb4e1..f4a2d648 100644 --- a/optree/ops.py +++ b/optree/ops.py @@ -25,16 +25,7 @@ import sys import textwrap from collections import OrderedDict, defaultdict, deque -from typing import ( - TYPE_CHECKING, - Any, - Callable, - ClassVar, - Generic, - SupportsIndex, - SupportsInt, - overload, -) +from typing import TYPE_CHECKING, Any, ClassVar, Generic, SupportsIndex, SupportsInt, overload import optree._C as _C from optree.accessors import PyTreeAccessor @@ -43,7 +34,7 @@ if TYPE_CHECKING: - from collections.abc import Collection, Iterable, Mapping + from collections.abc import Callable, Collection, Iterable, Mapping from optree.accessors import PyTreeEntry from optree.typing import ( @@ -1249,7 +1240,7 @@ def tree_transpose( leaves[offset : offset + inner_size] for offset in range(0, outer_size * inner_size, inner_size) ] - transposed = zip(*grouped) + transposed = zip(*grouped, strict=True) subtrees = map(outer_treespec.unflatten, transposed) return inner_treespec.unflatten(subtrees) # type: ignore[arg-type] @@ -1353,7 +1344,9 @@ def tree_transpose_map( raise ValueError(f'The inner structure must have at least one leaf. Got: {inner_treespec}.') grouped = [inner_treespec.flatten_up_to(o) for o in outputs] - transposed = zip(*grouped) if grouped else [()] * inner_treespec.num_leaves + # Every group is `inner_treespec.flatten_up_to(...)`, hence the same length. `zip()` over no + # groups yields nothing at all, so a leafless outer structure spells its empty subtrees out. + transposed = zip(*grouped, strict=True) if grouped else [()] * inner_treespec.num_leaves subtrees = map(outer_treespec.unflatten, transposed) return inner_treespec.unflatten(subtrees) # type: ignore[arg-type] @@ -1440,7 +1433,9 @@ def tree_transpose_map_with_path( raise ValueError(f'The inner structure must have at least one leaf. Got: {inner_treespec}.') grouped = [inner_treespec.flatten_up_to(o) for o in outputs] - transposed = zip(*grouped) if grouped else [()] * inner_treespec.num_leaves + # Every group is `inner_treespec.flatten_up_to(...)`, hence the same length. `zip()` over no + # groups yields nothing at all, so a leafless outer structure spells its empty subtrees out. + transposed = zip(*grouped, strict=True) if grouped else [()] * inner_treespec.num_leaves subtrees = map(outer_treespec.unflatten, transposed) return inner_treespec.unflatten(subtrees) # type: ignore[arg-type] @@ -1554,7 +1549,9 @@ def tree_transpose_map_with_accessor( raise ValueError(f'The inner structure must have at least one leaf. Got: {inner_treespec}.') grouped = [inner_treespec.flatten_up_to(o) for o in outputs] - transposed = zip(*grouped) if grouped else [()] * inner_treespec.num_leaves + # Every group is `inner_treespec.flatten_up_to(...)`, hence the same length. `zip()` over no + # groups yields nothing at all, so a leafless outer structure spells its empty subtrees out. + transposed = zip(*grouped, strict=True) if grouped else [()] * inner_treespec.num_leaves subtrees = map(outer_treespec.unflatten, transposed) return inner_treespec.unflatten(subtrees) # type: ignore[arg-type] @@ -2845,12 +2842,15 @@ def treespec_transform( >>> treespec = tree_structure({'b': 3, 'a': (0, [1, 2]), 'c': (4, None)}) >>> treespec PyTreeSpec({'a': (*, [*, *]), 'b': *, 'c': (*, None)}) - >>> treespec_transform(treespec, lambda spec: treespec_dict(zip(spec.entries(), spec.children()))) + >>> treespec_transform( + ... treespec, + ... lambda spec: treespec_dict(zip(spec.entries(), spec.children(), strict=True)), + ... ) PyTreeSpec({'a': {0: *, 1: {0: *, 1: *}}, 'b': *, 'c': {0: *, 1: {}}}) >>> treespec_transform( ... treespec, ... lambda spec: ( - ... treespec_ordereddict(zip(spec.entries(), spec.children())) + ... treespec_ordereddict(zip(spec.entries(), spec.children(), strict=True)) ... if spec.type is dict ... else spec ... ), @@ -3875,7 +3875,7 @@ def helper( # pylint: disable=too-many-locals for e in prefix_tree_entries ] # pylint: disable-next=invalid-name - for e, t1, t2 in zip(entries, prefix_tree_children, full_tree_children): + for e, t1, t2 in zip(entries, prefix_tree_children, full_tree_children, strict=True): yield from helper(accessor + e, t1, t2) # type: ignore[arg-type] return list(helper(PyTreeAccessor(), prefix_tree, full_tree)) diff --git a/optree/pytree.py b/optree/pytree.py index 3e725652..4e711bdb 100644 --- a/optree/pytree.py +++ b/optree/pytree.py @@ -132,8 +132,7 @@ if _TYPE_CHECKING: from collections.abc import Callable, Iterable - from typing import Any, TypeVar # pylint: disable=ungrouped-imports - from typing_extensions import ParamSpec # Python 3.10+ + from typing import Any, ParamSpec, TypeVar # pylint: disable=ungrouped-imports _P = ParamSpec('_P') _T = TypeVar('_T') diff --git a/optree/registry.py b/optree/registry.py index 340fd212..e2292909 100644 --- a/optree/registry.py +++ b/optree/registry.py @@ -27,7 +27,7 @@ from collections import OrderedDict, defaultdict, deque, namedtuple from operator import itemgetter, methodcaller from threading import Lock -from typing import TYPE_CHECKING, Any, Callable, ClassVar, Generic, NamedTuple, TypeVar, overload +from typing import TYPE_CHECKING, Any, ClassVar, Generic, NamedTuple, TypeVar, overload import optree._C as _C from optree.accessors import ( @@ -52,7 +52,7 @@ if TYPE_CHECKING: - from collections.abc import Collection, Generator, Iterable + from collections.abc import Callable, Collection, Generator, Iterable from optree.typing import VT, CustomTreeNode, FlattenFunc, UnflattenFunc @@ -68,12 +68,11 @@ ] -SLOTS = {'slots': True} if sys.version_info >= (3, 10) else {} # Python 3.10+ if sys.version_info >= (3, 15) and _C.OPTREE_HAS_FROZENDICT: # pragma: >=3.15 cover from builtins import frozendict # pylint: disable=no-name-in-module -@dataclasses.dataclass(init=True, repr=True, eq=True, frozen=True, **SLOTS) +@dataclasses.dataclass(init=True, repr=True, eq=True, frozen=True, slots=True) class PyTreeNodeRegistryEntry(Generic[T]): """A dataclass that stores the information of a pytree node type.""" @@ -81,17 +80,13 @@ class PyTreeNodeRegistryEntry(Generic[T]): flatten_func: FlattenFunc[T] unflatten_func: UnflattenFunc[T] - if sys.version_info >= (3, 10): # pragma: >=3.10 cover - _: dataclasses.KW_ONLY # Python 3.10+ + _: dataclasses.KW_ONLY path_entry_type: builtins.type[PyTreeEntry] = AutoEntry kind: PyTreeKind = PyTreeKind.CUSTOM namespace: str = '' -del SLOTS - - # pylint: disable-next=missing-class-docstring,too-few-public-methods class GlobalNamespace: # pragma: no cover __slots__: ClassVar[tuple[()]] = () @@ -106,7 +101,7 @@ def __repr__(self, /) -> str: if TYPE_CHECKING: - from typing_extensions import ParamSpec # Python 3.10+ + from typing import ParamSpec _P = ParamSpec('_P') _T = TypeVar('_T') @@ -222,11 +217,9 @@ def pytree_node_registry_get( # noqa: C901 and cls is not namedtuple # noqa: PYI024 and not inspect.isclass(cls) ): - raise TypeError(f'Expected a class or None, got {cls!r}.') # pragma: !=3.9 cover + raise TypeError(f'Expected a class or None, got {cls!r}.') if not isinstance(namespace, str): - raise TypeError( # pragma: !=3.9 cover - f'The namespace must be a string, got {namespace!r}.', - ) + raise TypeError(f'The namespace must be a string, got {namespace!r}.') if cls is None: with __REGISTRY_LOCK: @@ -361,7 +354,9 @@ def register_pytree_node( ... list(vars(ct).keys()), ... list(vars(ct).keys()), ... ), - ... unflatten_func=lambda keys, values: MyContainer(**dict(zip(keys, values))), + ... unflatten_func=lambda keys, values: MyContainer( + ... **dict(zip(keys, values, strict=True)), + ... ), ... path_entry_type=MyContainerEntry, ... namespace='mycontainer', ... ) diff --git a/optree/typing.py b/optree/typing.py index 42c9d118..9905c581 100644 --- a/optree/typing.py +++ b/optree/typing.py @@ -17,6 +17,7 @@ from __future__ import annotations import abc +import atexit import functools import platform import sys @@ -29,6 +30,7 @@ from collections import defaultdict as DefaultDict # noqa: N812 from collections import deque as Deque # noqa: N812 from collections.abc import ( + Callable, Collection, Hashable, ItemsView, @@ -40,13 +42,13 @@ ) from typing import ( Any, - Callable, ClassVar, Final, ForwardRef, Generic, - Optional, + ParamSpec, Protocol, + TypeAlias, TypeVar, Union, final, @@ -56,9 +58,7 @@ from typing_extensions import ( NamedTuple, # Generic NamedTuple: Python 3.11+ Never, # Python 3.11+ - ParamSpec, # Python 3.10+ Self, # Python 3.11+ - TypeAlias, # Python 3.10+ TypeAliasType, # Python 3.12+ ) from weakref import WeakKeyDictionary @@ -149,7 +149,7 @@ Children: TypeAlias = Iterable[T] -MetaData: TypeAlias = Optional[Hashable] +MetaData: TypeAlias = Hashable | None @runtime_checkable @@ -173,7 +173,12 @@ def __tree_unflatten__(cls, metadata: MetaData, children: Children[T], /) -> Sel """Unflatten the children and metadata into the custom pytree node.""" -_UnionType = type(Union[int, str]) +# Before Python 3.14, `Union[int, str]` produces `typing._UnionGenericAlias` while `int | str` +# produces `types.UnionType` -- they are different types. On Python 3.14+, the two are unified and +# `Union[int, str]` also produces `types.UnionType`. Using `type(Union[int, str])` here ensures +# `_UnionType` automatically matches the pytree alias type on all supported Python versions. See +# the comment at `__class_getitem__` below for why the pytree aliases use `Union[...]`. +_UnionType = type(Union[int, str]) # noqa: UP007 try: # pragma: no cover @@ -183,6 +188,8 @@ def __tree_unflatten__(cls, metadata: MetaData, children: Children[T], /) -> Sel def _tp_cache(func: Callable[P, T], /) -> Callable[P, T]: cached = functools.lru_cache(func) + atexit.register(cached.cache_clear) + @functools.wraps(func) def inner(*args: P.args, **kwargs: P.kwargs) -> T: try: @@ -274,6 +281,11 @@ def __class_getitem__( # noqa: C901 # pylint: disable=too-many-branches else: recurse_ref = ForwardRef(f'{cls.__name__}[{param!r}]') + # We use `Union[...]` explicitly rather than chained `|` for clarity. Before Python 3.14, + # chained `|` with `typing._GenericAlias` operands (e.g., `Tuple[x]`, `List[y]`) would still + # produce `typing._UnionGenericAlias` (not `types.UnionType`) via `__or__`/`__ror__`. + # On Python 3.14+, both `Union[...]` and `|` produce `types.UnionType`. + # TODO(PEP 604): migrate to `|` when minimum Python is raised to 3.14+. pytree_types = [ param, Tuple[recurse_ref, ...], # type: ignore[valid-type] # Tuple, NamedTuple, PyStructSequence @@ -288,7 +300,7 @@ def __class_getitem__( # noqa: C901 # pylint: disable=too-many-branches CustomTreeNode[recurse_ref], # type: ignore[list-item,valid-type] ], ) - pytree_alias = Union[tuple(pytree_types)] # type: ignore[valid-type] + pytree_alias = Union[tuple(pytree_types)] # type: ignore[valid-type] # noqa: UP007 with cls.__instance_lock__: cls.__instances__[pytree_alias] = (param, name) # type: ignore[index] @@ -628,7 +640,9 @@ def structseq_fields(obj: tuple | type[tuple], /) -> tuple[str, ...]: positions = [index_of[id(getattr(probe, name))] for name in named] except (TypeError, ValueError, KeyError, AttributeError): # pragma: no cover pass # the type rejects placeholder values, fall back to assuming a trailing layout - names_by_index = dict(zip(positions, named)) + # `strict=False`: on the fallback above, `positions` spans every slot while `named` covers + # only the named ones, and truncating to the shorter is the trailing-layout assumption. + names_by_index = dict(zip(positions, named, strict=False)) return tuple( names_by_index.get(index, PyStructSequence_UnnamedField) diff --git a/optree/utils.py b/optree/utils.py index 1878757e..49f725a2 100644 --- a/optree/utils.py +++ b/optree/utils.py @@ -16,9 +16,9 @@ from __future__ import annotations -from collections.abc import Iterable, Sequence +from collections.abc import Callable, Iterable, Sequence from operator import itemgetter -from typing import TYPE_CHECKING, Any, Callable, overload +from typing import TYPE_CHECKING, Any, overload if TYPE_CHECKING: @@ -41,7 +41,7 @@ def total_order_sorted( # Apply `key` up front: it runs exactly once per element, and a `TypeError` from the callback # propagates instead of being mistaken for a comparison failure and swallowed below. keys: list[Any] = sequence if key is None else [key(x) for x in sequence] - decorated = list(zip(keys, sequence)) + decorated = list(zip(keys, sequence, strict=True)) def by_type_and_key(pair: tuple[Any, T], /) -> tuple[str, Any]: y = pair[0] @@ -104,7 +104,7 @@ def safe_zip(*args: Iterable[Any]) -> zip[tuple[Any, ...]]: seqs = [arg if isinstance(arg, Sequence) else list(arg) for arg in args] if len(set(map(len, seqs))) > 1: raise ValueError(f'length mismatch: {list(map(len, seqs))}') - return zip(*seqs) + return zip(*seqs, strict=True) def unzip2(xys: Iterable[tuple[T, S]], /) -> tuple[tuple[T, ...], tuple[S, ...]]: diff --git a/pyproject.toml b/pyproject.toml index 8dc63bad..f31e0c3f 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -8,7 +8,7 @@ build-backend = "setuptools.build_meta" name = "optree" description = "Optimized PyTree Utilities." readme = "README.md" -requires-python = ">= 3.9" +requires-python = ">= 3.10" authors = [{ name = "OpTree Contributors" }] license = "Apache-2.0" keywords = [ @@ -21,7 +21,6 @@ classifiers = [ "Development Status :: 4 - Beta", "Programming Language :: C++", "Programming Language :: Python :: 3", - "Programming Language :: Python :: 3.9", "Programming Language :: Python :: 3.10", "Programming Language :: Python :: 3.11", "Programming Language :: Python :: 3.12", @@ -165,7 +164,7 @@ environment = { CMAKE_OSX_SYSROOT = "iphonesimulator" } # Linter tools ################################################################# [tool.mypy] -python_version = "3.9" +python_version = "3.10" exclude = ['^tests/.*\.py$', '^(third|3rd)[_\-]?party/.*$', '^\.?venv/.*$'] pretty = true show_column_numbers = true @@ -188,7 +187,7 @@ warn_unused_ignores = true no_site_packages = true [tool.pylint] -main.py-version = "3.9" +main.py-version = "3.10" main.extension-pkg-allow-list = ["optree._C"] main.ignore-paths = ['^tests/$', '^(third|3rd)[_\-]?party/$'] basic.good-names = [] @@ -214,7 +213,7 @@ builtin = "clear,rare,en-GB_to_en-US" ignore-words = "docs/source/spelling_wordlist.txt" [tool.ruff] -target-version = "py39" +target-version = "py310" line-length = 100 output-format = "full" src = ["optree", "tests"] diff --git a/tests/concurrent/test_subinterpreters.py b/tests/concurrent/test_subinterpreters.py index ece31031..bca7dae1 100644 --- a/tests/concurrent/test_subinterpreters.py +++ b/tests/concurrent/test_subinterpreters.py @@ -449,6 +449,7 @@ def test_registry_init_failure_does_not_leak_the_interpreter_id(): textwrap.dedent( ''' import atexit + import typing def failing_register(*args, **kwargs): raise RuntimeError('injected atexit failure') diff --git a/tests/helpers.py b/tests/helpers.py index bb22ff3b..e7ef708d 100644 --- a/tests/helpers.py +++ b/tests/helpers.py @@ -132,7 +132,7 @@ def represent(value): return repr(value) ids = tuple( - '-'.join(f'{arg}({represent(value)})' for arg, value in zip(arguments, values)) + '-'.join(f'{arg}({represent(value)})' for arg, value in zip(arguments, values, strict=True)) for values in argvalues ) @@ -186,7 +186,7 @@ def check_script_in_subprocess( /, *, output, - timeout=120.0, + timeout=300.0, cwd=TEST_ROOT, env=None, rerun=1, @@ -241,7 +241,7 @@ def assert_equal_type_and_value(actual, expected=MISSING, *, expected_type=None) assert actual == expected if isinstance(expected, optree.PyTreeAccessor): assert hash(actual) == hash(expected) - for i, j in zip(actual, expected): + for i, j in zip(actual, expected, strict=True): assert_equal_type_and_value(i, j) @@ -478,7 +478,7 @@ def __tree_flatten__(self): @classmethod def __tree_unflatten__(cls, metadata, children): - return cls(zip(metadata, children)) + return cls(zip(metadata, children, strict=True)) def __repr__(self): return f'{self.__class__.__name__}({super().__repr__()})' diff --git a/tests/integrations/test_jax.py b/tests/integrations/test_jax.py index aee46bb6..c0efe244 100644 --- a/tests/integrations/test_jax.py +++ b/tests/integrations/test_jax.py @@ -72,7 +72,7 @@ def replace_leaf(_): reconstructed_leaves, reconstructed_treespec = optree.tree_flatten(reconstructed) assert reconstructed_treespec == treespec assert len(leaves) == len(reconstructed_leaves) - for leaf, reconstructed_leaf in zip(leaves, reconstructed_leaves): + for leaf, reconstructed_leaf in zip(leaves, reconstructed_leaves, strict=True): assert jnp.allclose(leaf, reconstructed_leaf) leaf = jnp.asarray(leaf) reconstructed_leaf = jnp.asarray(reconstructed_leaf) @@ -133,7 +133,7 @@ def replace_leaf(_): reconstructed_leaves, reconstructed_treespec = optree.tree_flatten(reconstructed) assert reconstructed_treespec == treespec assert len(leaves) == len(reconstructed_leaves) - for leaf, reconstructed_leaf in zip(leaves, reconstructed_leaves): + for leaf, reconstructed_leaf in zip(leaves, reconstructed_leaves, strict=True): assert jnp.allclose(leaf, reconstructed_leaf) leaf = jnp.asarray(leaf) reconstructed_leaf = jnp.asarray(reconstructed_leaf) diff --git a/tests/integrations/test_numpy.py b/tests/integrations/test_numpy.py index 7da5d297..985a622d 100644 --- a/tests/integrations/test_numpy.py +++ b/tests/integrations/test_numpy.py @@ -66,7 +66,7 @@ def replace_leaf(_): reconstructed_leaves, reconstructed_treespec = optree.tree_flatten(reconstructed) assert reconstructed_treespec == treespec assert len(leaves) == len(reconstructed_leaves) - for leaf, reconstructed_leaf in zip(leaves, reconstructed_leaves): + for leaf, reconstructed_leaf in zip(leaves, reconstructed_leaves, strict=True): assert np.allclose(leaf, reconstructed_leaf) leaf = np.asarray(leaf) reconstructed_leaf = np.asarray(reconstructed_leaf) @@ -126,7 +126,7 @@ def replace_leaf(_): reconstructed_leaves, reconstructed_treespec = optree.tree_flatten(reconstructed) assert reconstructed_treespec == treespec assert len(leaves) == len(reconstructed_leaves) - for leaf, reconstructed_leaf in zip(leaves, reconstructed_leaves): + for leaf, reconstructed_leaf in zip(leaves, reconstructed_leaves, strict=True): assert np.allclose(leaf, reconstructed_leaf) leaf = np.asarray(leaf) reconstructed_leaf = np.asarray(reconstructed_leaf) diff --git a/tests/integrations/test_torch.py b/tests/integrations/test_torch.py index eccfc2de..bd86cce2 100644 --- a/tests/integrations/test_torch.py +++ b/tests/integrations/test_torch.py @@ -74,7 +74,7 @@ def replace_leaf(_): reconstructed_leaves, reconstructed_treespec = optree.tree_flatten(reconstructed) assert reconstructed_treespec == treespec assert len(leaves) == len(reconstructed_leaves) - for leaf, reconstructed_leaf in zip(leaves, reconstructed_leaves): + for leaf, reconstructed_leaf in zip(leaves, reconstructed_leaves, strict=True): assert torch.is_tensor(leaf) assert torch.is_tensor(reconstructed_leaf) assert torch.allclose(leaf, reconstructed_leaf) @@ -136,7 +136,7 @@ def replace_leaf(_): reconstructed_leaves, reconstructed_treespec = optree.tree_flatten(reconstructed) assert reconstructed_treespec == treespec assert len(leaves) == len(reconstructed_leaves) - for leaf, reconstructed_leaf in zip(leaves, reconstructed_leaves): + for leaf, reconstructed_leaf in zip(leaves, reconstructed_leaves, strict=True): assert torch.is_tensor(leaf) assert torch.is_tensor(reconstructed_leaf) assert torch.allclose(leaf, reconstructed_leaf) diff --git a/tests/test_accessors.py b/tests/test_accessors.py index cb1def40..40e9f034 100644 --- a/tests/test_accessors.py +++ b/tests/test_accessors.py @@ -583,7 +583,7 @@ def __tree_unflatten__(cls, metadata, children): accessors, leaves, _ = optree.tree_flatten_with_accessor(obj, namespace='namespace') assert leaves == [1, 2, 3] assert accessors == expected_accessors - for a, b in zip(accessors, expected_accessors): + for a, b in zip(accessors, expected_accessors, strict=True): assert_equal_type_and_value(a, b) for accessor in accessors: diff --git a/tests/test_dataclasses.py b/tests/test_dataclasses.py index 5d53802a..9458387d 100644 --- a/tests/test_dataclasses.py +++ b/tests/test_dataclasses.py @@ -106,20 +106,10 @@ def test_field_future_parameters(): optree.dataclasses.field() dataclasses.field() - if sys.version_info >= (3, 10): - optree.dataclasses.field(kw_only=True) - dataclasses.field(kw_only=True) - optree.dataclasses.field(kw_only=False) - dataclasses.field(kw_only=False) - else: - with pytest.raises(TypeError, match='got an unexpected keyword argument'): - optree.dataclasses.field(kw_only=True) - with pytest.raises(TypeError, match='got an unexpected keyword argument'): - dataclasses.field(kw_only=True) - with pytest.raises(TypeError, match='got an unexpected keyword argument'): - optree.dataclasses.field(kw_only=False) - with pytest.raises(TypeError, match='got an unexpected keyword argument'): - dataclasses.field(kw_only=False) + optree.dataclasses.field(kw_only=True) + dataclasses.field(kw_only=True) + optree.dataclasses.field(kw_only=False) + dataclasses.field(kw_only=False) if sys.version_info >= (3, 14): optree.dataclasses.field(doc='doc') @@ -179,41 +169,18 @@ def test_dataclass_future_parameters(): optree.dataclasses.dataclass(namespace='namespace') dataclasses.dataclass() - if sys.version_info >= (3, 10): - optree.dataclasses.dataclass(match_args=True, namespace='namespace') - dataclasses.dataclass(match_args=True) - optree.dataclasses.dataclass(match_args=False, namespace='namespace') - dataclasses.dataclass(match_args=False) - optree.dataclasses.dataclass(kw_only=True, namespace='namespace') - dataclasses.dataclass(kw_only=True) - optree.dataclasses.dataclass(kw_only=False, namespace='namespace') - dataclasses.dataclass(kw_only=False) - optree.dataclasses.dataclass(slots=True, namespace='namespace') - dataclasses.dataclass(slots=True) - optree.dataclasses.dataclass(slots=False, namespace='namespace') - dataclasses.dataclass(slots=False) - else: - optree.dataclasses.dataclass(match_args=True, namespace='namespace') - with pytest.raises(TypeError, match='got an unexpected keyword argument'): - dataclasses.dataclass(match_args=True) - with pytest.raises(TypeError, match='got an unexpected keyword argument'): - optree.dataclasses.dataclass(match_args=False, namespace='error') - with pytest.raises(TypeError, match='got an unexpected keyword argument'): - dataclasses.dataclass(match_args=False) - with pytest.raises(TypeError, match='got an unexpected keyword argument'): - optree.dataclasses.dataclass(kw_only=True, namespace='error') - with pytest.raises(TypeError, match='got an unexpected keyword argument'): - dataclasses.dataclass(kw_only=True) - optree.dataclasses.dataclass(kw_only=False, namespace='namespace') - with pytest.raises(TypeError, match='got an unexpected keyword argument'): - dataclasses.dataclass(kw_only=False) - with pytest.raises(TypeError, match='got an unexpected keyword argument'): - optree.dataclasses.dataclass(slots=True, namespace='error') - with pytest.raises(TypeError, match='got an unexpected keyword argument'): - dataclasses.dataclass(slots=True) - optree.dataclasses.dataclass(slots=False, namespace='namespace') - with pytest.raises(TypeError, match='got an unexpected keyword argument'): - dataclasses.dataclass(slots=False) + optree.dataclasses.dataclass(match_args=True, namespace='namespace') + dataclasses.dataclass(match_args=True) + optree.dataclasses.dataclass(match_args=False, namespace='namespace') + dataclasses.dataclass(match_args=False) + optree.dataclasses.dataclass(kw_only=True, namespace='namespace') + dataclasses.dataclass(kw_only=True) + optree.dataclasses.dataclass(kw_only=False, namespace='namespace') + dataclasses.dataclass(kw_only=False) + optree.dataclasses.dataclass(slots=True, namespace='namespace') + dataclasses.dataclass(slots=True) + optree.dataclasses.dataclass(slots=False, namespace='namespace') + dataclasses.dataclass(slots=False) if sys.version_info >= (3, 11): optree.dataclasses.dataclass(weakref_slot=True, namespace='namespace') @@ -442,149 +409,72 @@ def test_make_dataclass_future_parameters(): }, ) - if sys.version_info >= (3, 10): - optree.dataclasses.make_dataclass( - 'Foo2', - ['x', ('y', int), ('z', float, 0.0)], - match_args=True, - namespace='namespace', - ) - dataclasses.make_dataclass( - 'Foo3', - ['x', ('y', int), ('z', float, 0.0)], - match_args=True, - ) - optree.dataclasses.make_dataclass( - 'Foo4', - ['x', ('y', int), ('z', float, 0.0)], - match_args=False, - namespace='namespace', - ) - dataclasses.make_dataclass( - 'Foo5', - ['x', ('y', int), ('z', float, 0.0)], - match_args=False, - ) - optree.dataclasses.make_dataclass( - 'Foo6', - ['x', ('y', int), ('z', float, 0.0)], - kw_only=True, - namespace='namespace', - ) - dataclasses.make_dataclass( - 'Foo7', - ['x', ('y', int), ('z', float, 0.0)], - kw_only=True, - ) - optree.dataclasses.make_dataclass( - 'Foo8', - ['x', ('y', int), ('z', float, 0.0)], - kw_only=False, - namespace='namespace', - ) - dataclasses.make_dataclass( - 'Foo9', - ['x', ('y', int), ('z', float, 0.0)], - kw_only=False, - ) - optree.dataclasses.make_dataclass( - 'Foo10', - ['x', ('y', int), ('z', float, 0.0)], - slots=True, - namespace='namespace', - ) - dataclasses.make_dataclass( - 'Foo11', - ['x', ('y', int), ('z', float, 0.0)], - slots=True, - ) - optree.dataclasses.make_dataclass( - 'Foo12', - ['x', ('y', int), ('z', float, 0.0)], - slots=False, - namespace='namespace', - ) - dataclasses.make_dataclass( - 'Foo13', - ['x', ('y', int), ('z', float, 0.0)], - slots=False, - ) - else: - optree.dataclasses.make_dataclass( - 'Foo2', - ['x', ('y', int), ('z', float, 0.0)], - match_args=True, - namespace='namespace', - ) - with pytest.raises(TypeError, match='got an unexpected keyword argument'): - dataclasses.make_dataclass( - 'Foo3', - ['x', ('y', int), ('z', float, 0.0)], - match_args=True, - ) - with pytest.raises(TypeError, match='got an unexpected keyword argument'): - optree.dataclasses.make_dataclass( - 'Foo4', - ['x', ('y', int), ('z', float, 0.0)], - match_args=False, - namespace='error', - ) - with pytest.raises(TypeError, match='got an unexpected keyword argument'): - dataclasses.make_dataclass( - 'Foo5', - ['x', ('y', int), ('z', float, 0.0)], - match_args=False, - ) - with pytest.raises(TypeError, match='got an unexpected keyword argument'): - optree.dataclasses.make_dataclass( - 'Foo6', - ['x', ('y', int), ('z', float, 0.0)], - kw_only=True, - namespace='error', - ) - with pytest.raises(TypeError, match='got an unexpected keyword argument'): - dataclasses.make_dataclass( - 'Foo7', - ['x', ('y', int), ('z', float, 0.0)], - kw_only=True, - ) - optree.dataclasses.make_dataclass( - 'Foo8', - ['x', ('y', int), ('z', float, 0.0)], - kw_only=False, - namespace='namespace', - ) - with pytest.raises(TypeError, match='got an unexpected keyword argument'): - dataclasses.make_dataclass( - 'Foo9', - ['x', ('y', int), ('z', float, 0.0)], - kw_only=False, - ) - with pytest.raises(TypeError, match='got an unexpected keyword argument'): - optree.dataclasses.make_dataclass( - 'Foo10', - ['x', ('y', int), ('z', float, 0.0)], - slots=True, - namespace='error', - ) - with pytest.raises(TypeError, match='got an unexpected keyword argument'): - dataclasses.make_dataclass( - 'Foo11', - ['x', ('y', int), ('z', float, 0.0)], - slots=True, - ) - optree.dataclasses.make_dataclass( - 'Foo12', - ['x', ('y', int), ('z', float, 0.0)], - slots=False, - namespace='namespace', - ) - with pytest.raises(TypeError, match='got an unexpected keyword argument'): - dataclasses.make_dataclass( - 'Foo13', - ['x', ('y', int), ('z', float, 0.0)], - slots=False, - ) + optree.dataclasses.make_dataclass( + 'Foo2', + ['x', ('y', int), ('z', float, 0.0)], + match_args=True, + namespace='namespace', + ) + dataclasses.make_dataclass( + 'Foo3', + ['x', ('y', int), ('z', float, 0.0)], + match_args=True, + ) + optree.dataclasses.make_dataclass( + 'Foo4', + ['x', ('y', int), ('z', float, 0.0)], + match_args=False, + namespace='namespace', + ) + dataclasses.make_dataclass( + 'Foo5', + ['x', ('y', int), ('z', float, 0.0)], + match_args=False, + ) + optree.dataclasses.make_dataclass( + 'Foo6', + ['x', ('y', int), ('z', float, 0.0)], + kw_only=True, + namespace='namespace', + ) + dataclasses.make_dataclass( + 'Foo7', + ['x', ('y', int), ('z', float, 0.0)], + kw_only=True, + ) + optree.dataclasses.make_dataclass( + 'Foo8', + ['x', ('y', int), ('z', float, 0.0)], + kw_only=False, + namespace='namespace', + ) + dataclasses.make_dataclass( + 'Foo9', + ['x', ('y', int), ('z', float, 0.0)], + kw_only=False, + ) + optree.dataclasses.make_dataclass( + 'Foo10', + ['x', ('y', int), ('z', float, 0.0)], + slots=True, + namespace='namespace', + ) + dataclasses.make_dataclass( + 'Foo11', + ['x', ('y', int), ('z', float, 0.0)], + slots=True, + ) + optree.dataclasses.make_dataclass( + 'Foo12', + ['x', ('y', int), ('z', float, 0.0)], + slots=False, + namespace='namespace', + ) + dataclasses.make_dataclass( + 'Foo13', + ['x', ('y', int), ('z', float, 0.0)], + slots=False, + ) if sys.version_info >= (3, 11): optree.dataclasses.make_dataclass( diff --git a/tests/test_ops.py b/tests/test_ops.py index f4ab1aee..439407d3 100644 --- a/tests/test_ops.py +++ b/tests/test_ops.py @@ -21,7 +21,6 @@ import itertools import operator import pickle -import platform import re import sys from collections import OrderedDict, defaultdict, deque @@ -41,7 +40,6 @@ CustomTuple, FlatCache, MyAnotherDict, - Py_DEBUG, always, assert_equal_type_and_value, check_script_in_subprocess, @@ -65,9 +63,6 @@ def test_import_no_warnings(): def test_max_depth(): - if sys.version_info < (3, 10) and platform.system() == 'Windows' and Py_DEBUG: - pytest.skip('Flaky with Python 3.9 on Windows in debug mode.') - lst = [1] for _ in range(optree.MAX_RECURSION_DEPTH - 1): lst = [lst] @@ -705,7 +700,7 @@ def test_paths_and_accessors(data): assert other_treespec == expected_treespec assert paths == expected_paths assert accessors == expected_accessors - for leaf, accessor, path in zip(leaves, accessors, paths): + for leaf, accessor, path in zip(leaves, accessors, paths, strict=True): assert isinstance(accessor, optree.PyTreeAccessor) assert isinstance(path, tuple) assert len(accessor) == len(path) @@ -776,7 +771,7 @@ def test_paths_and_accessors_with_is_leaf( assert treespec == expected_treespec assert other_leaves == expected_leaves assert other_treespec == expected_treespec - for leaf, accessor, path in zip(leaves, accessors, paths): + for leaf, accessor, path in zip(leaves, accessors, paths, strict=True): assert isinstance(accessor, optree.PyTreeAccessor) assert isinstance(path, tuple) assert len(accessor) == len(path) @@ -1564,6 +1559,93 @@ def test_tree_transpose_map_with_accessor(): ) +@parametrize( + transpose_map=[ + optree.tree_transpose_map, + optree.tree_transpose_map_with_path, + optree.tree_transpose_map_with_accessor, + ], +) +def test_tree_transpose_map_with_no_leaves(transpose_map): + # The zero-leaf guard belongs to inferring the inner structure from the first output. With an + # explicit `inner_treespec` there is nothing to infer, so a leafless outer structure transposes + # into one copy of itself per inner leaf and `func` is never called. Those empty subtrees must + # be spelled out: `zip(*[])` yields nothing, and `unflatten` would fail with `Too few leaves`. + # `test_tree_partition_with_no_leaves` covers the same branch via `tree_transpose_map` only. + def must_not_be_called(*args): + raise AssertionError(f'`func` must not be called for a leafless tree, got {args!r}.') + + for tree in [(), [], {}, {'a': [], 'b': ()}, [(), {}], {'a': None}]: + out = transpose_map( + must_not_be_called, + tree, + inner_treespec=optree.tree_structure({'x': 0, 'y': 1}), + ) + assert out == {'x': tree, 'y': tree}, (tree, out) + + # One empty subtree per inner LEAF, not one per inner child: a nested inner structure gets a + # copy of the outer structure at every leaf position. + out = transpose_map( + must_not_be_called, + tree, + inner_treespec=optree.tree_structure({'x': 0, 'y': (1, 2)}), + ) + assert out == {'x': tree, 'y': (tree, tree)}, (tree, out) + + # `rests` are flattened up to the (leafless) outer structure, so they contribute no outputs + # either. + out = transpose_map( + must_not_be_called, + tree, + tree, + inner_treespec=optree.tree_structure({'x': 0, 'y': 1}), + ) + assert out == {'x': tree, 'y': tree}, (tree, out) + + # Without an explicit `inner_treespec` the inner structure is inferred from the first output, so + # the outer structure must still have at least one leaf. + with pytest.raises( + ValueError, + match=r'The outer structure must have at least one leaf\. Got: .*\.', + ): + transpose_map(must_not_be_called, ()) + + +@parametrize( + transpose_map=[ + optree.tree_transpose_map, + optree.tree_transpose_map_with_path, + optree.tree_transpose_map_with_accessor, + ], +) +def test_tree_transpose_map_flatten_up_to_is_rectangular(transpose_map): + # Pins the premise behind `zip(..., strict=True)` rather than the flag, which by design never + # fires and so cannot be caught by any test. The groups are `inner_treespec.flatten_up_to(o)`, + # which returns exactly `inner_treespec.num_leaves` items or raises, never a short list. So all + # groups have one length whatever shape the outputs take, and a structure mismatch is reported + # by `flatten_up_to`, naming the offending output, rather than by the `zip()`. + inner_treespec = optree.tree_structure({'x': 0, 'y': 1}) + assert [ + len(inner_treespec.flatten_up_to(output)) + for output in ({'x': 0, 'y': 1}, {'x': (0, 1), 'y': [2]}, {'x': None, 'y': {'z': 3}}) + ] == [inner_treespec.num_leaves] * 3 + + def make_output(*args): + # `tree_transpose_map` passes the leaf alone, while the path and accessor variants prepend + # one argument, so the leaf is the last one either way. + leaf = args[-1] + if leaf % 2: + return {'x': (leaf, leaf), 'y': [leaf]} + return {'x': leaf, 'y': leaf} + + # The outputs above have different shapes per leaf, yet each flattens up to 2 subtrees. + out = transpose_map(make_output, [1, 2, 3], inner_treespec=inner_treespec) + assert out == {'x': [(1, 1), 2, (3, 3)], 'y': [[1], 2, [3]]} + + with pytest.raises(ValueError, match=re.escape('dictionary key mismatch')): + transpose_map(lambda *args: {'x': 0}, [1, 2], inner_treespec=inner_treespec) + + def test_tree_map_none_is_leaf(): x = ((1, 2, None), (3, 4, 5)) y = (([6], None, None), ({'foo': 'bar'}, 7, [8, 9])) @@ -3509,7 +3591,7 @@ def test_tree_broadcast_map_unchanged_for_structural_predicates(): # The common structure still honors `is_leaf` where it describes the INPUT trees, and plain # broadcasting is unaffected by the change. assert optree.tree_broadcast_map( - lambda x, y: x + y, + operator.add, [1, 2], [[3, 4], [5, 6]], ) == [[4, 5], [7, 8]] @@ -3875,7 +3957,7 @@ def flatten(node): # noqa: C901 assert node_kind == optree.PyTreeKind.CUSTOM assert len(entries) == len(children) if hasattr(node, '__getitem__'): - for child, entry in zip(children, entries): + for child, entry in zip(children, entries, strict=True): assert node[entry] is child if node_type is type(None): @@ -3883,7 +3965,7 @@ def flatten(node): # noqa: C901 with pytest.raises(ValueError, match=re.escape('Expected no children.')): unflatten_func(metadata, range(1)) - for child, entry in zip(children, entries): + for child, entry in zip(children, entries, strict=True): path_stack.append(entry) accessor_stack.append(output.path_entry_type(entry, node_type, node_kind)) flatten(child) diff --git a/tests/test_registry.py b/tests/test_registry.py index 9df76afc..4486c1a4 100644 --- a/tests/test_registry.py +++ b/tests/test_registry.py @@ -33,7 +33,6 @@ GLOBAL_NAMESPACE, NODETYPE_REGISTRY, OPTREE_HAS_FROZENDICT, - PYPY, Py_GIL_DISABLED, check_script_in_subprocess, disable_systrace, @@ -1002,11 +1001,10 @@ def test_pytree_node_registry_get_with_invalid_arguments(): assert optree.register_pytree_node.get(None) == registry assert optree.register_pytree_node.get(namespace=GLOBAL_NAMESPACE) == registry assert optree.register_pytree_node.get(namedtuple) is registry[namedtuple] # noqa: PYI024 - if sys.version_info[:2] != (3, 9) or PYPY: - with pytest.raises(TypeError, match='Expected a class or None'): - optree.register_pytree_node.get(dataclass) - with pytest.raises(TypeError, match='The namespace must be a string'): - optree.register_pytree_node.get(list, namespace=None) + with pytest.raises(TypeError, match='Expected a class or None'): + optree.register_pytree_node.get(dataclass) + with pytest.raises(TypeError, match='The namespace must be a string'): + optree.register_pytree_node.get(list, namespace=None) def test_pytree_node_registry_with_init_subclass(): @@ -1022,7 +1020,7 @@ def __tree_flatten__(self): @classmethod def __tree_unflatten__(cls, metadata, children): - return cls(zip(metadata, children)) + return cls(zip(metadata, children, strict=True)) class MyAnotherDict(MyDict): pass diff --git a/tests/test_treespec.py b/tests/test_treespec.py index 3d14ae70..e266c8d2 100644 --- a/tests/test_treespec.py +++ b/tests/test_treespec.py @@ -553,6 +553,7 @@ class Shared: @pytest.mark.xfail( strict=True, reason='known limitation: a treespec cannot see registration references held by another treespec', + raises=AssertionError, ) def test_treespec_reference_cycle_across_treespecs_is_collectable(): # Known limitation. A treespec reports a registration's members only when its own nodes hold @@ -1452,7 +1453,11 @@ def test_treespec_compose_children( stack = [(composed_treespec.children(), expected_treespec.children())] while stack: composed_children, expected_children = stack.pop() - for composed_child, expected_child in zip(composed_children, expected_children): + for composed_child, expected_child in zip( + composed_children, + expected_children, + strict=True, + ): assert composed_child == expected_child stack.append((composed_child.children(), expected_child.children())) @@ -2063,7 +2068,7 @@ def gen_path(spec): yield () return - for entry, child in zip(entries, children): + for entry, child in zip(entries, children, strict=True): for suffix in gen_path(child): yield (entry, *suffix) @@ -2092,7 +2097,7 @@ def gen_typed_path(spec): node_type = spec.type node_kind = spec.kind - for entry, child in zip(entries, children): + for entry, child in zip(entries, children, strict=True): for suffix in gen_typed_path(child): yield ((entry, node_type, node_kind), *suffix) @@ -2375,11 +2380,11 @@ def test_treespec_transform(): ) == optree.tree_structure([[1, 2, 3], [4]]) assert optree.treespec_transform( treespec, - lambda spec: optree.treespec_dict(zip('abcd', spec.children())), + lambda spec: optree.treespec_dict(zip('abcd', spec.children(), strict=False)), ) == optree.tree_structure({'a': {'a': 0, 'b': 1, 'c': 2}, 'b': {'a': 3}}) assert optree.treespec_transform( treespec, - lambda spec: optree.treespec_dict(zip('abcd', spec.children())), + lambda spec: optree.treespec_dict(zip('abcd', spec.children(), strict=False)), lambda spec: optree.tree_structure([0, None, 1]), ) == optree.tree_structure( {'a': {'a': [0, None, 1], 'b': [2, None, 3], 'c': [4, None, 5]}, 'b': {'a': [6, None, 7]}}, @@ -2392,7 +2397,7 @@ def test_treespec_transform(): optree.treespec_transform( treespec, lambda spec: optree.tree_structure( - MyAnotherDict(zip(spec.entries(), spec.children())), + MyAnotherDict(zip(spec.entries(), spec.children(), strict=True)), namespace='namespace', ), ) @@ -2432,7 +2437,10 @@ def test_treespec_transform(): def fn(spec): with optree.dict_insertion_ordered(True, namespace='undefined'): - return optree.treespec_dict(zip('abcd', spec.children()), namespace='undefined') + return optree.treespec_dict( + zip('abcd', spec.children(), strict=False), + namespace='undefined', + ) with pytest.raises(ValueError, match=r'Expected treespec\(s\) with namespace .*, got .*\.'): optree.treespec_transform(namespaced_treespec, fn) @@ -2884,7 +2892,7 @@ def test_treespec_constructor( # noqa: C901 if use_sorted_keys: assert ( optree.treespec_dict( - zip(sorted(node), children_treespecs), + zip(sorted(node), children_treespecs, strict=True), none_is_leaf=none_is_leaf, namespace=passed_namespace, ) @@ -2892,7 +2900,7 @@ def test_treespec_constructor( # noqa: C901 ) assert ( optree.treespec_from_collection( - dict(zip(sorted(node), children_treespecs)), + dict(zip(sorted(node), children_treespecs, strict=True)), none_is_leaf=none_is_leaf, namespace=passed_namespace, ) @@ -2910,7 +2918,7 @@ def test_treespec_constructor( # noqa: C901 with context: assert ( optree.treespec_dict( - zip(node, children_treespecs), + zip(node, children_treespecs, strict=True), none_is_leaf=none_is_leaf, namespace=passed_namespace, ) @@ -2918,7 +2926,7 @@ def test_treespec_constructor( # noqa: C901 ) assert ( optree.treespec_from_collection( - dict(zip(node, children_treespecs)), + dict(zip(node, children_treespecs, strict=True)), none_is_leaf=none_is_leaf, namespace=passed_namespace, ) @@ -2927,7 +2935,7 @@ def test_treespec_constructor( # noqa: C901 elif node_type is OrderedDict: assert ( optree.treespec_ordereddict( - zip(node, children_treespecs), + zip(node, children_treespecs, strict=True), none_is_leaf=none_is_leaf, namespace=passed_namespace, ) @@ -2935,7 +2943,7 @@ def test_treespec_constructor( # noqa: C901 ) assert ( optree.treespec_from_collection( - OrderedDict(zip(node, children_treespecs)), + OrderedDict(zip(node, children_treespecs, strict=True)), none_is_leaf=none_is_leaf, namespace=passed_namespace, ) @@ -2946,7 +2954,7 @@ def test_treespec_constructor( # noqa: C901 assert ( optree.treespec_defaultdict( node.default_factory, - zip(sorted(node), children_treespecs), + zip(sorted(node), children_treespecs, strict=True), none_is_leaf=none_is_leaf, namespace=passed_namespace, ) @@ -2956,7 +2964,7 @@ def test_treespec_constructor( # noqa: C901 optree.treespec_from_collection( defaultdict( node.default_factory, - zip(sorted(node), children_treespecs), + zip(sorted(node), children_treespecs, strict=True), ), none_is_leaf=none_is_leaf, namespace=passed_namespace, @@ -2976,7 +2984,7 @@ def test_treespec_constructor( # noqa: C901 assert ( optree.treespec_defaultdict( node.default_factory, - zip(node, children_treespecs), + zip(node, children_treespecs, strict=True), none_is_leaf=none_is_leaf, namespace=passed_namespace, ) @@ -2986,7 +2994,7 @@ def test_treespec_constructor( # noqa: C901 optree.treespec_from_collection( defaultdict( node.default_factory, - zip(node, children_treespecs), + zip(node, children_treespecs, strict=True), ), none_is_leaf=none_is_leaf, namespace=passed_namespace, @@ -3001,7 +3009,7 @@ def test_treespec_constructor( # noqa: C901 if use_sorted_keys: assert ( optree.treespec_frozendict( - zip(sorted(node), children_treespecs), + zip(sorted(node), children_treespecs, strict=True), none_is_leaf=none_is_leaf, namespace=passed_namespace, ) @@ -3010,7 +3018,7 @@ def test_treespec_constructor( # noqa: C901 assert ( optree.treespec_from_collection( builtins.frozendict( # type: ignore[attr-defined] - zip(sorted(node), children_treespecs), + zip(sorted(node), children_treespecs, strict=True), ), none_is_leaf=none_is_leaf, namespace=passed_namespace, @@ -3029,7 +3037,7 @@ def test_treespec_constructor( # noqa: C901 with context: assert ( optree.treespec_frozendict( - zip(node, children_treespecs), + zip(node, children_treespecs, strict=True), none_is_leaf=none_is_leaf, namespace=passed_namespace, ) @@ -3038,7 +3046,7 @@ def test_treespec_constructor( # noqa: C901 assert ( optree.treespec_from_collection( builtins.frozendict( # type: ignore[attr-defined] - zip(node, children_treespecs), + zip(node, children_treespecs, strict=True), ), none_is_leaf=none_is_leaf, namespace=passed_namespace, diff --git a/tests/test_typing.py b/tests/test_typing.py index 24c757a6..25d26f0f 100644 --- a/tests/test_typing.py +++ b/tests/test_typing.py @@ -28,6 +28,7 @@ import pytest import optree +import optree.typing from helpers import ( OPTREE_HAS_FROZENDICT, PYBIND11_HAS_NATIVE_ENUM, @@ -121,14 +122,13 @@ def test_pytree_typing(): T = TypeVar('T') optree.PyTree[int] - optree.PyTree[Union[int, str]] + optree.PyTree[Union[int, str]] # noqa: UP007 optree.PyTree[T] assert optree.PyTree[optree.PyTree[int]] == optree.PyTree[int] - assert optree.PyTree[optree.PyTree[Union[int, str]]] == optree.PyTree[Union[int, str]] + assert optree.PyTree[optree.PyTree[Union[int, str]]] == optree.PyTree[Union[int, str]] # noqa: UP007 assert optree.PyTree[optree.PyTree[T]] == optree.PyTree[T] - if sys.version_info >= (3, 10): - optree.PyTree[int | str] - assert optree.PyTree[optree.PyTree[int | str]] == optree.PyTree[int | str] + optree.PyTree[float | bytes] + assert optree.PyTree[optree.PyTree[float | bytes]] == optree.PyTree[float | bytes] IntTree = optree.PyTreeTypeVar('IntTree', int) # noqa: N806 assert IntTree == optree.PyTree[IntTree] @@ -647,6 +647,45 @@ def test_structseq_fields(): structseq_fields(FakeStructSequence) +@skipif_pypy # PyPy reports `n_unnamed_fields == 0` and takes the index-based branch instead +def test_structseq_fields_python_implementation_falls_back_when_the_probe_is_rejected( + monkeypatch: pytest.MonkeyPatch, +): + # A type with unnamed slots that rejects the sentinel probe leaves nothing to match positions + # against, so the implementation falls back to a trailing layout: the named fields keep the + # leading positions and the rest get the unnamed marker. Pinning the lenient pairing that makes + # that possible, since a strict one would raise instead of degrading. + # + # No stdlib type reaches this path: `os.stat_result`, the only CPython type with unnamed fields, + # accepts arbitrary objects. The stand-in below is not a real PyStructSequence, hence the + # `is_structseq_class` swap; `__slots__` supplies the `member_descriptor` fields it looks for. + class RejectsProbe: + __slots__ = ('st_alpha', 'st_beta', 'st_gamma') + + n_fields = 4 + n_sequence_fields = 4 + n_unnamed_fields = 1 + + def __new__(cls, sequence, /): + raise TypeError(f'cannot build {cls.__name__} from placeholder values: {sequence!r}') + + python_implementation = optree.structseq_fields.__python_implementation__ + original_is_structseq_class = optree.typing.is_structseq_class + monkeypatch.setattr( + optree.typing, + 'is_structseq_class', + lambda cls, /: cls is RejectsProbe or original_is_structseq_class(cls), + ) + fields = python_implementation(RejectsProbe) + + assert fields == ( + 'st_alpha', + 'st_beta', + 'st_gamma', + optree.typing.PyStructSequence_UnnamedField, + ) + + def test_structseq_accessor_unnamed_fields_codify_by_index(): # The accessor round-trip (the generated code evaluates to the accessed value) must hold for # every slot on every implementation. It exercises both codify styles: CPython leaves @@ -756,6 +795,7 @@ def test_type_caches_register_interpreter_cleanup(): r""" import atexit import time + import typing n0 = atexit._ncallbacks() import optree @@ -839,6 +879,7 @@ def test_type_cache_insert_failure_before_import_does_not_crash(): r""" import atexit import sys + import typing real_register = atexit.register