refactor!: mark some arguments as positional-only

optree/_C.pyi

@@ -17,12 +17,11 @@
 
 import builtins
 import enum
-from collections.abc import Callable, Iterable, Iterator
+from collections.abc import Callable, Collection, Iterable, Iterator
 from typing import Any
 from typing_extensions import Self
 
 from optree.typing import (
-    CustomTreeNode,
     FlattenFunc,
     MetaData,
     PyTree,
@@ -44,12 +43,14 @@ GLIBCXX_USE_CXX11_ABI: bool
 
 def flatten(
     tree: PyTree[T],
+    /,
     leaf_predicate: Callable[[T], bool] | None = None,
     node_is_leaf: bool = False,
     namespace: str = '',
 ) -> tuple[list[T], PyTreeSpec]: ...
 def flatten_with_path(
     tree: PyTree[T],
+    /,
     leaf_predicate: Callable[[T], bool] | None = None,
     node_is_leaf: bool = False,
     namespace: str = '',
@@ -63,30 +64,33 @@ def make_none(
     namespace: str = '',  # unused
 ) -> PyTreeSpec: ...
 def make_from_collection(
-    collection: CustomTreeNode[PyTreeSpec],
+    collection: Collection[PyTreeSpec],
+    /,
     node_is_leaf: bool = False,
     namespace: str = '',
 ) -> PyTreeSpec: ...
 def is_leaf(
     obj: T,
+    /,
     leaf_predicate: Callable[[T], bool] | None = None,
     node_is_leaf: bool = False,
     namespace: str = '',
 ) -> bool: ...
 def all_leaves(
     iterable: Iterable[T],
+    /,
     leaf_predicate: Callable[[T], bool] | None = None,
     node_is_leaf: bool = False,
     namespace: str = '',
 ) -> bool: ...
-def is_namedtuple(obj: object | type) -> bool: ...
-def is_namedtuple_instance(obj: object) -> bool: ...
-def is_namedtuple_class(cls: type) -> bool: ...
-def namedtuple_fields(obj: tuple | type[tuple]) -> tuple[str, ...]: ...
-def is_structseq(obj: object | type) -> bool: ...
-def is_structseq_instance(obj: object) -> bool: ...
-def is_structseq_class(cls: type) -> bool: ...
-def structseq_fields(obj: tuple | type[tuple]) -> tuple[str, ...]: ...
+def is_namedtuple(obj: object | type, /) -> bool: ...
+def is_namedtuple_instance(obj: object, /) -> bool: ...
+def is_namedtuple_class(cls: type, /) -> bool: ...
+def namedtuple_fields(obj: tuple | type[tuple], /) -> tuple[str, ...]: ...
+def is_structseq(obj: object | type, /) -> bool: ...
+def is_structseq_instance(obj: object, /) -> bool: ...
+def is_structseq_class(cls: type, /) -> bool: ...
+def structseq_fields(obj: tuple | type[tuple], /) -> tuple[str, ...]: ...
 
 class PyTreeKind(enum.IntEnum):
     CUSTOM = 0  # a custom type
@@ -109,10 +113,10 @@ class PyTreeSpec:
     namespace: str
     type: builtins.type | None
     kind: PyTreeKind
-    def unflatten(self, leaves: Iterable[T]) -> PyTree[T]: ...
-    def flatten_up_to(self, full_tree: PyTree[T]) -> list[PyTree[T]]: ...
-    def broadcast_to_common_suffix(self, other: PyTreeSpec) -> PyTreeSpec: ...
-    def compose(self, inner_treespec: PyTreeSpec) -> PyTreeSpec: ...
+    def unflatten(self, leaves: Iterable[T], /) -> PyTree[T]: ...
+    def flatten_up_to(self, full_tree: PyTree[T], /) -> list[PyTree[T]]: ...
+    def broadcast_to_common_suffix(self, other: PyTreeSpec, /) -> PyTreeSpec: ...
+    def compose(self, inner_treespec: PyTreeSpec, /) -> PyTreeSpec: ...
     def walk(
         self,
         f_node: Callable[[tuple[U, ...], MetaData], U],
@@ -122,25 +126,26 @@ class PyTreeSpec:
     def paths(self) -> list[tuple[Any, ...]]: ...
     def accessors(self) -> list[PyTreeAccessor]: ...
     def entries(self) -> list[Any]: ...
-    def entry(self, index: int) -> Any: ...
+    def entry(self, index: int, /) -> Any: ...
     def children(self) -> list[PyTreeSpec]: ...
-    def child(self, index: int) -> PyTreeSpec: ...
+    def child(self, index: int, /) -> PyTreeSpec: ...
     def is_leaf(self, strict: bool = True) -> bool: ...
-    def is_prefix(self, other: PyTreeSpec, strict: bool = False) -> bool: ...
-    def is_suffix(self, other: PyTreeSpec, strict: bool = False) -> bool: ...
-    def __eq__(self, other: object) -> bool: ...
-    def __ne__(self, other: object) -> bool: ...
-    def __lt__(self, other: object) -> bool: ...
-    def __le__(self, other: object) -> bool: ...
-    def __gt__(self, other: object) -> bool: ...
-    def __ge__(self, other: object) -> bool: ...
+    def is_prefix(self, other: PyTreeSpec, /, strict: bool = False) -> bool: ...
+    def is_suffix(self, other: PyTreeSpec, /, strict: bool = False) -> bool: ...
+    def __eq__(self, other: object, /) -> bool: ...
+    def __ne__(self, other: object, /) -> bool: ...
+    def __lt__(self, other: object, /) -> bool: ...
+    def __le__(self, other: object, /) -> bool: ...
+    def __gt__(self, other: object, /) -> bool: ...
+    def __ge__(self, other: object, /) -> bool: ...
     def __hash__(self) -> int: ...
     def __len__(self) -> int: ...
 
 class PyTreeIter(Iterator[T]):
     def __init__(
         self,
         tree: PyTree[T],
+        /,
         leaf_predicate: Callable[[T], bool] | None = None,
         node_is_leaf: bool = False,
         namespace: str = '',
@@ -149,14 +154,16 @@ class PyTreeIter(Iterator[T]):
     def __next__(self) -> T: ...
 
 def register_node(
-    cls: type[CustomTreeNode[T]],
-    flatten_func: FlattenFunc,
-    unflatten_func: UnflattenFunc,
+    cls: type[Collection[T]],
+    /,
+    flatten_func: FlattenFunc[T],
+    unflatten_func: UnflattenFunc[T],
     path_entry_type: type[PyTreeEntry],
     namespace: str = '',
 ) -> None: ...
 def unregister_node(
-    cls: type[CustomTreeNode[T]],
+    cls: type,
+    /,
     namespace: str = '',
 ) -> None: ...
 def is_dict_insertion_ordered(
@@ -165,5 +172,6 @@ def is_dict_insertion_ordered(
 ) -> bool: ...
 def set_dict_insertion_ordered(
     mode: bool,
+    /,
     namespace: str = '',
 ) -> None: ...

optree/accessor.py

@@ -76,7 +76,7 @@ def __post_init__(self) -> None:
         if self.kind == PyTreeKind.NONE:
             raise ValueError('Cannot create a path entry for None.')
 
-    def __call__(self, obj: Any) -> Any:
+    def __call__(self, obj: Any, /) -> Any:
         """Get the child object."""
         try:
             return obj[self.entry]  # should be overridden
@@ -85,15 +85,15 @@ def __call__(self, obj: Any) -> Any:
                 f'{self.__class__!r} cannot access through {obj!r} via entry {self.entry!r}',
             ) from ex
 
-    def __add__(self, other: object) -> PyTreeAccessor:
+    def __add__(self, other: object, /) -> PyTreeAccessor:
         """Join the path entry with another path entry or accessor."""
         if isinstance(other, PyTreeEntry):
             return PyTreeAccessor((self, other))
         if isinstance(other, PyTreeAccessor):
             return PyTreeAccessor((self, *other))
         return NotImplemented
 
-    def __eq__(self, other: object) -> bool:
+    def __eq__(self, other: object, /) -> bool:
         """Check if the path entries are equal."""
         return isinstance(other, PyTreeEntry) and (
             (
@@ -196,7 +196,7 @@ class GetItemEntry(PyTreeEntry):
 
     __slots__: ClassVar[tuple[()]] = ()
 
-    def __call__(self, obj: Any) -> Any:
+    def __call__(self, obj: Any, /) -> Any:
         """Get the child object."""
         return obj[self.entry]
 
@@ -217,7 +217,7 @@ def name(self) -> str:
         """Get the attribute name."""
         return self.entry
 
-    def __call__(self, obj: Any) -> Any:
+    def __call__(self, obj: Any, /) -> Any:
         """Get the child object."""
         return getattr(obj, self.name)
 
@@ -245,7 +245,7 @@ def index(self) -> int:
         """Get the index."""
         return self.entry
 
-    def __call__(self, obj: Sequence[_T_co]) -> _T_co:
+    def __call__(self, obj: Sequence[_T_co], /) -> _T_co:
         """Get the child object."""
         return obj[self.index]
 
@@ -267,7 +267,7 @@ def key(self) -> _KT_co:
         """Get the key."""
         return self.entry
 
-    def __call__(self, obj: Mapping[_KT_co, _VT_co]) -> _VT_co:
+    def __call__(self, obj: Mapping[_KT_co, _VT_co], /) -> _VT_co:
         """Get the child object."""
         return obj[self.key]
 
@@ -383,43 +383,43 @@ def __new__(cls, path: Iterable[PyTreeEntry] = ()) -> Self:
             raise TypeError(f'Expected a path of PyTreeEntry, got {path!r}.')
         return super().__new__(cls, path)
 
-    def __call__(self, obj: Any) -> Any:
+    def __call__(self, obj: Any, /) -> Any:
         """Get the child object."""
         for entry in self:
             obj = entry(obj)
         return obj
 
     @overload  # type: ignore[override]
-    def __getitem__(self, index: int) -> PyTreeEntry:  # noqa: D105,RUF100
+    def __getitem__(self, index: int, /) -> PyTreeEntry:  # noqa: D105,RUF100
         ...
 
     @overload
-    def __getitem__(self, index: slice) -> PyTreeAccessor:  # noqa: D105,RUF100
+    def __getitem__(self, index: slice, /) -> PyTreeAccessor:  # noqa: D105,RUF100
         ...
 
-    def __getitem__(self, index: int | slice) -> PyTreeEntry | PyTreeAccessor:
+    def __getitem__(self, index: int | slice, /) -> PyTreeEntry | PyTreeAccessor:
         """Get the child path entry or an accessor for a subpath."""
         if isinstance(index, slice):
             return PyTreeAccessor(super().__getitem__(index))
         return super().__getitem__(index)
 
-    def __add__(self, other: object) -> PyTreeAccessor:
+    def __add__(self, other: object, /) -> PyTreeAccessor:
         """Join the accessor with another path entry or accessor."""
         if isinstance(other, PyTreeEntry):
             return PyTreeAccessor((*self, other))
         if isinstance(other, PyTreeAccessor):
             return PyTreeAccessor((*self, *other))
         return NotImplemented
 
-    def __mul__(self, value: int) -> PyTreeAccessor:  # type: ignore[override]
+    def __mul__(self, value: int, /) -> PyTreeAccessor:  # type: ignore[override]
         """Repeat the accessor."""
         return PyTreeAccessor(super().__mul__(value))
 
-    def __rmul__(self, value: int) -> PyTreeAccessor:  # type: ignore[override]
+    def __rmul__(self, value: int, /) -> PyTreeAccessor:  # type: ignore[override]
         """Repeat the accessor."""
         return PyTreeAccessor(super().__rmul__(value))
 
-    def __eq__(self, other: object) -> bool:
+    def __eq__(self, other: object, /) -> bool:
         """Check if the accessors are equal."""
         return isinstance(other, PyTreeAccessor) and super().__eq__(other)
 

optree/dataclasses.py

@@ -155,6 +155,7 @@ def field(  # type: ignore[no-redef] # pylint: disable=function-redefined,too-ma
 @dataclass_transform(field_specifiers=(field,))
 def dataclass(  # pylint: disable=too-many-arguments
     cls: None,
+    /,
     *,
     init: bool = True,
     repr: bool = True,  # pylint: disable=redefined-builtin
@@ -175,6 +176,7 @@ def dataclass(  # pylint: disable=too-many-arguments
 @dataclass_transform(field_specifiers=(field,))
 def dataclass(  # pylint: disable=too-many-arguments
     cls: _TypeT,
+    /,
     *,
     init: bool = True,
     repr: bool = True,  # pylint: disable=redefined-builtin
@@ -193,6 +195,7 @@ def dataclass(  # pylint: disable=too-many-arguments
 @dataclass_transform(field_specifiers=(field,))
 def dataclass(  # noqa: C901 # pylint: disable=function-redefined,too-many-arguments,too-many-locals,too-many-branches
     cls: _TypeT | None = None,
+    /,
     *,
     init: bool = True,
     repr: bool = True,  # pylint: disable=redefined-builtin

optree/functools.py

@@ -42,13 +42,13 @@ class _HashablePartialShim:
     args: tuple[Any, ...]
     keywords: dict[str, Any]
 
-    def __init__(self, partial_func: functools.partial) -> None:
+    def __init__(self, partial_func: functools.partial, /) -> None:
         self.partial_func: functools.partial = partial_func
 
     def __call__(self, *args: Any, **kwargs: Any) -> Any:
         return self.partial_func(*args, **kwargs)
 
-    def __eq__(self, other: object) -> bool:
+    def __eq__(self, other: object, /) -> bool:
         if isinstance(other, _HashablePartialShim):
             return self.partial_func == other.partial_func
         return self.partial_func == other
@@ -118,7 +118,7 @@ class partial(  # noqa: N801 # pylint: disable=invalid-name,too-few-public-metho
 
     TREE_PATH_ENTRY_TYPE: ClassVar[type[PyTreeEntry]] = GetAttrEntry
 
-    def __new__(cls, func: Callable[..., Any], *args: T, **keywords: T) -> Self:
+    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