feat(ops): add more info to tree_flatten_one_level (#168)

.github/workflows/tests-with-pydebug.yml

@@ -80,6 +80,9 @@ jobs:
       - name: Determine Python version
         shell: bash
         run: |
+          if [[ "${{ runner.os }}" == 'Windows' ]]; then
+            pyenv update
+          fi
           pyenv install --list
           if [[ "${{ matrix.python-abiflags }}" == *t* ]]; then
             PYTHON_VERSION="$(
@@ -111,6 +114,7 @@ jobs:
               PYENV_INSTALL_ARGS+=("--debug")
             fi
           else
+            pyenv update
             if [[ "${{ matrix.python-abiflags }}" == *d* ]]; then
               PYENV_INSTALL_ARGS+=("--dev")
             fi

CHANGELOG.md

@@ -13,6 +13,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ### Added
 
+- Add more info to `tree_flatten_one_level` by [@XuehaiPan](https://github.com/XuehaiPan) in [#168](https://github.com/metaopt/optree/pull/168).
 - Improve typing support for generic `PyTree[T]` and registry lookup / register functions by [@XuehaiPan](https://github.com/XuehaiPan) in [#160](https://github.com/metaopt/optree/pull/160) and [#166](https://github.com/metaopt/optree/pull/166).
 
 ### Changed

README.md

@@ -121,7 +121,7 @@ OpTree out-of-box supports the following Python container types in the registry:
 
 which are considered non-leaf nodes in the tree.
 Python objects that the type is not registered will be treated as leaf nodes.
-The registration lookup uses the `is` operator to determine whether the type is matched.
+The registry lookup uses the `is` operator to determine whether the type is matched.
 So subclasses will need to explicitly register in the registry, otherwise, an object of that type will be considered a leaf.
 The [`NoneType`](https://docs.python.org/3/library/constants.html#None) is a special case discussed in section [`None` is non-leaf Node vs. `None` is Leaf](#none-is-non-leaf-node-vs-none-is-leaf).
 
@@ -289,7 +289,7 @@ There are several key attributes of the pytree type registry:
 
 3. **Users cannot modify the behavior of already registered built-in types** listed in [Built-in PyTree Node Types](#built-in-pytree-node-types), such as key order sorting for `dict` and `collections.defaultdict`.
 
-4. **Inherited subclasses are not implicitly registered.** The registration lookup uses `type(obj) is registered_type` rather than `isinstance(obj, registered_type)`. Users need to register the subclasses explicitly. To register all subclasses, it is easy to implement with [`metaclass`](https://docs.python.org/3/reference/datamodel.html#metaclasses) or [`__init_subclass__`](https://docs.python.org/3/reference/datamodel.html#customizing-class-creation), for example:
+4. **Inherited subclasses are not implicitly registered.** The registry lookup uses `type(obj) is registered_type` rather than `isinstance(obj, registered_type)`. Users need to register the subclasses explicitly. To register all subclasses, it is easy to implement with [`metaclass`](https://docs.python.org/3/reference/datamodel.html#metaclasses) or [`__init_subclass__`](https://docs.python.org/3/reference/datamodel.html#customizing-class-creation), for example:
 
     ```python
     from collections import UserDict

docs/source/spelling_wordlist.txt

@@ -94,3 +94,6 @@ hypot
 init
 ns
 metaclass
+cpu
+classmethod
+unregister

optree/accessor.py

@@ -344,7 +344,7 @@ class DataclassEntry(GetAttrEntry):
     entry: str | int  # type: ignore[assignment]
 
     @property
-    def fields(self) -> tuple[str, ...]:
+    def fields(self) -> tuple[str, ...]:  # pragma: no cover
         """Get all field names."""
         return tuple(f.name for f in dataclasses.fields(self.type))
 

optree/dataclasses.py

@@ -68,9 +68,6 @@
 from typing_extensions import Literal  # Python 3.8+
 from typing_extensions import dataclass_transform  # Python 3.11+
 
-from optree.accessor import DataclassEntry
-from optree.registry import register_pytree_node
-
 
 # Redefine `field`, `dataclasses`, and `make_dataclasses`.
 # The remaining APIs are re-exported from the original package.
@@ -338,9 +335,12 @@ def unflatten_func(metadata: tuple[tuple[str, Any], ...], children: tuple[_U, ..
         kwargs.update(metadata)
         return cls(**kwargs)
 
+    from optree.accessor import DataclassEntry  # pylint: disable=import-outside-toplevel
+    from optree.registry import register_pytree_node  # pylint: disable=import-outside-toplevel
+
     return register_pytree_node(  # type: ignore[return-value]
         cls,
-        flatten_func,
+        flatten_func,  # type: ignore[arg-type]
         unflatten_func,  # type: ignore[arg-type]
         path_entry_type=DataclassEntry,
         namespace=namespace,

optree/functools.py

@@ -17,16 +17,20 @@
 from __future__ import annotations
 
 import functools
-from typing import Any, Callable, ClassVar
+from typing import TYPE_CHECKING, Any, Callable, ClassVar
 from typing_extensions import Self  # Python 3.11+
 from typing_extensions import deprecated  # Python 3.13+
 
 from optree import registry
-from optree.accessor import GetAttrEntry, PyTreeEntry
+from optree.accessor import GetAttrEntry
 from optree.ops import tree_reduce as reduce
 from optree.typing import CustomTreeNode, T
 
 
+if TYPE_CHECKING:
+    from optree.accessor import PyTreeEntry
+
+
 __all__ = [
     'partial',
     'reduce',