Skip to content

Commit

Permalink
refactor!: mark some arguments as positional-only
Browse files Browse the repository at this point in the history
  • Loading branch information
@XuehaiPan
XuehaiPan committed Oct 25, 2024
1 parent 52df8f9 commit 3603954
Show file tree
Hide file tree
Showing 10 changed files with 274 additions and 182 deletions.
72 changes: 41 additions & 31 deletions optree/_C.pyi
View file
Original file line number Diff line number Diff line change
Expand Up @@ -43,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 = '',
Expand All @@ -63,29 +65,32 @@ def make_none(
) -> PyTreeSpec: ...
def make_from_collection(
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
Expand All @@ -108,54 +113,58 @@ 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: Self, /) -> Self: ...
def compose(self, inner_treespec: Self, /) -> Self: ...
def walk(
self,
/,
f_node: Callable[[tuple[U, ...], MetaData], U],
f_leaf: Callable[[T], U] | None,
leaves: Iterable[T],
) -> U: ...
def paths(self) -> list[tuple[Any, ...]]: ...
def accessors(self) -> list[PyTreeAccessor]: ...
def entries(self) -> list[Any]: ...
def entry(self, index: int) -> Any: ...
def children(self) -> list[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 __hash__(self) -> int: ...
def __len__(self) -> int: ...
def paths(self, /) -> list[tuple[Any, ...]]: ...
def accessors(self, /) -> list[PyTreeAccessor]: ...
def entries(self, /) -> list[Any]: ...
def entry(self, index: int, /) -> Any: ...
def children(self, /) -> list[Self]: ...
def child(self, index: int, /) -> Self: ...
def is_leaf(self, /, strict: bool = True) -> bool: ...
def is_prefix(self, other: Self, /, strict: bool = False) -> bool: ...
def is_suffix(self, other: Self, /, 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 = '',
) -> None: ...
def __iter__(self) -> Self: ...
def __next__(self) -> T: ...
def __iter__(self, /) -> Self: ...
def __next__(self, /) -> T: ...

def register_node(
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,
/,
namespace: str = '',
) -> None: ...
def is_dict_insertion_ordered(
Expand All @@ -164,5 +173,6 @@ def is_dict_insertion_ordered(
) -> bool: ...
def set_dict_insertion_ordered(
mode: bool,
/,
namespace: str = '',
) -> None: ...
30 changes: 15 additions & 15 deletions optree/accessor.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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 (
(
Expand Down Expand Up @@ -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]

Expand All @@ -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)

Expand Down Expand Up @@ -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]

Expand All @@ -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]

Expand Down Expand Up @@ -388,43 +388,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)

Expand Down
20 changes: 12 additions & 8 deletions optree/dataclasses.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -56,6 +56,8 @@
True
"""

# pylint: disable=too-many-arguments

from __future__ import annotations

import contextlib
Expand Down Expand Up @@ -86,7 +88,7 @@


@overload # type: ignore[no-redef]
def field( # pylint: disable=too-many-arguments
def field(
*,
default: _T,
init: bool = True,
Expand All @@ -101,7 +103,7 @@ def field( # pylint: disable=too-many-arguments


@overload
def field( # pylint: disable=too-many-arguments
def field(
*,
default_factory: Callable[[], _T],
init: bool = True,
Expand All @@ -115,7 +117,7 @@ def field( # pylint: disable=too-many-arguments


@overload
def field( # pylint: disable=too-many-arguments
def field(
*,
init: bool = True,
repr: bool = True, # pylint: disable=redefined-builtin
Expand All @@ -127,7 +129,7 @@ def field( # pylint: disable=too-many-arguments
) -> Any: ...


def field( # pylint: disable=function-redefined,too-many-arguments
def field( # pylint: disable=function-redefined
*,
default: Any = dataclasses.MISSING,
default_factory: Any = dataclasses.MISSING,
Expand Down Expand Up @@ -194,7 +196,7 @@ def field( # pylint: disable=function-redefined,too-many-arguments


@overload # type: ignore[no-redef]
def dataclass( # pylint: disable=too-many-arguments
def dataclass(
*,
init: bool = True,
repr: bool = True, # pylint: disable=redefined-builtin
Expand All @@ -212,8 +214,9 @@ def dataclass( # pylint: disable=too-many-arguments


@overload
def dataclass( # pylint: disable=too-many-arguments
def dataclass(
cls: _TypeT,
/,
*,
init: bool = True,
repr: bool = True, # pylint: disable=redefined-builtin
Expand All @@ -230,8 +233,9 @@ 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
def dataclass( # noqa: C901 # pylint: disable=function-redefined,too-many-locals,too-many-branches
cls: _TypeT | None = None,
/,
*,
init: bool = True,
repr: bool = True, # pylint: disable=redefined-builtin
Expand Down Expand Up @@ -346,7 +350,7 @@ def unflatten_func(metadata: tuple[tuple[str, Any], ...], children: tuple[_U, ..
)


# pylint: disable-next=function-redefined,too-many-arguments,too-many-locals,too-many-branches
# pylint: disable-next=function-redefined,too-many-locals,too-many-branches
def make_dataclass( # type: ignore[no-redef] # noqa: C901
cls_name: str,
# pylint: disable-next=redefined-outer-name
Expand Down
6 changes: 3 additions & 3 deletions optree/functools.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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
Expand Down
Loading

0 comments on commit 3603954

Please sign in to comment.