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 committed Oct 9, 2024
1 parent c55c33b commit 6b5974a
Show file tree
Hide file tree
Showing 10 changed files with 255 additions and 164 deletions.
56 changes: 33 additions & 23 deletions optree/_C.pyi
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,38 +113,40 @@ 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],
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 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 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: ...

class PyTreeIter(Iterator[T]):
def __init__(
self,
tree: PyTree[T],
/,
leaf_predicate: Callable[[T], bool] | None = None,
node_is_leaf: bool = False,
namespace: str = '',
Expand All @@ -149,13 +156,15 @@ class PyTreeIter(Iterator[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
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 @@ -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)

Expand Down
3 changes: 3 additions & 0 deletions optree/dataclasses.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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
Expand All @@ -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
Expand Down
6 changes: 3 additions & 3 deletions optree/functools.py
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 6b5974a

Please sign in to comment.