1717from __future__ import annotations
1818
1919import dataclasses
20- import sys
2120from collections .abc import Iterable , Mapping , Sequence
2221from typing import TYPE_CHECKING , Any , ClassVar , Generic , Literal , TypeVar , overload
2322from typing_extensions import Self # Python 3.11+
4746]
4847
4948
50- SLOTS = {'slots' : True } if sys .version_info >= (3 , 10 ) else {} # Python 3.10+
51-
52-
53- @dataclasses .dataclass (init = True , repr = False , eq = False , frozen = True , ** SLOTS )
49+ @dataclasses .dataclass (init = True , repr = False , eq = False , frozen = True , slots = True )
5450class PyTreeEntry :
5551 """Base class for path entries."""
5652
@@ -122,9 +118,6 @@ def codify(self, /, node: str = '') -> str:
122118 return f'{ node } [<flat index { self .entry !r} >]' # should be overridden
123119
124120
125- del SLOTS
126-
127-
128121_T = TypeVar ('_T' )
129122_T_co = TypeVar ('_T_co' , covariant = True )
130123_KT_co = TypeVar ('_KT_co' , covariant = True )
@@ -134,8 +127,6 @@ def codify(self, /, node: str = '') -> str:
134127class AutoEntry (PyTreeEntry ):
135128 """A generic path entry class that determines the entry type on creation automatically."""
136129
137- __slots__ : ClassVar [tuple [()]] = ()
138-
139130 def __new__ ( # type: ignore[misc]
140131 cls ,
141132 / ,
@@ -184,8 +175,6 @@ def __new__( # type: ignore[misc]
184175class GetItemEntry (PyTreeEntry ):
185176 """A generic path entry class for nodes that access their children by :meth:`__getitem__`."""
186177
187- __slots__ : ClassVar [tuple [()]] = ()
188-
189178 def __call__ (self , obj : Any , / ) -> Any :
190179 """Get the child object."""
191180 return obj [self .entry ]
@@ -198,8 +187,6 @@ def codify(self, /, node: str = '') -> str:
198187class GetAttrEntry (PyTreeEntry ):
199188 """A generic path entry class for nodes that access their children by :meth:`__getattr__`."""
200189
201- __slots__ : ClassVar [tuple [()]] = ()
202-
203190 entry : str
204191
205192 @property
@@ -219,14 +206,10 @@ def codify(self, /, node: str = '') -> str:
219206class FlattenedEntry (PyTreeEntry ): # pylint: disable=too-few-public-methods
220207 """A fallback path entry class for flattened objects."""
221208
222- __slots__ : ClassVar [tuple [()]] = ()
223-
224209
225210class SequenceEntry (GetItemEntry , Generic [_T_co ]):
226211 """A path entry class for sequences."""
227212
228- __slots__ : ClassVar [tuple [()]] = ()
229-
230213 entry : int
231214 type : builtins .type [Sequence [_T_co ]]
232215
@@ -247,8 +230,6 @@ def __repr__(self, /) -> str:
247230class MappingEntry (GetItemEntry , Generic [_KT_co , _VT_co ]):
248231 """A path entry class for mappings."""
249232
250- __slots__ : ClassVar [tuple [()]] = ()
251-
252233 entry : _KT_co
253234 type : builtins .type [Mapping [_KT_co , _VT_co ]]
254235
@@ -269,8 +250,6 @@ def __repr__(self, /) -> str:
269250class NamedTupleEntry (SequenceEntry [_T ]):
270251 """A path entry class for namedtuple objects."""
271252
272- __slots__ : ClassVar [tuple [()]] = ()
273-
274253 entry : int
275254 type : builtins .type [NamedTuple [_T ]] # type: ignore[type-arg]
276255 kind : Literal [PyTreeKind .NAMEDTUPLE ]
@@ -299,8 +278,6 @@ def codify(self, /, node: str = '') -> str:
299278class StructSequenceEntry (SequenceEntry [_T ]):
300279 """A path entry class for PyStructSequence objects."""
301280
302- __slots__ : ClassVar [tuple [()]] = ()
303-
304281 entry : int
305282 type : builtins .type [StructSequence [_T ]]
306283 kind : Literal [PyTreeKind .STRUCTSEQUENCE ]
@@ -329,8 +306,6 @@ def codify(self, /, node: str = '') -> str:
329306class DataclassEntry (GetAttrEntry ):
330307 """A path entry class for dataclasses."""
331308
332- __slots__ : ClassVar [tuple [()]] = ()
333-
334309 entry : str | int # type: ignore[assignment]
335310
336311 @property
0 commit comments