@@ -73,9 +73,6 @@ def __init__(self, datapipe: IterDataPipe, key_value_fn: Optional[Callable] = No
73
73
self ._depleted = False
74
74
75
75
def _load_map (self ):
76
- if self ._map is None :
77
- self ._map = {}
78
- self ._itr = iter (self .datapipe )
79
76
while not self ._depleted :
80
77
try :
81
78
self ._load_next_item ()
@@ -84,10 +81,7 @@ def _load_map(self):
84
81
85
82
def __getitem__ (self , index ):
86
83
try :
87
- if self ._map is None :
88
- self ._map = {}
89
- self ._itr = iter (self .datapipe )
90
- else :
84
+ if self ._map is not None :
91
85
return self ._map [index ]
92
86
except KeyError :
93
87
pass
@@ -101,7 +95,10 @@ def __getitem__(self, index):
101
95
raise IndexError (f"Index { index } is invalid for IterToMapConverter." )
102
96
103
97
def _load_next_item (self ):
104
- elem = next (self ._itr )
98
+ if self ._map is None :
99
+ self ._map = {}
100
+ self ._itr = iter (self .datapipe )
101
+ elem = next (self ._itr ) # type: ignore[arg-type]
105
102
inp = elem if self .key_value_fn is None else self .key_value_fn (elem )
106
103
try :
107
104
length = len (inp )
@@ -135,14 +132,10 @@ def __getstate__(self):
135
132
dill_key_value_fn = dill .dumps (self .key_value_fn )
136
133
else :
137
134
dill_key_value_fn = self .key_value_fn
138
- return (
139
- self .datapipe ,
140
- dill_key_value_fn ,
141
- self ._map ,
142
- )
135
+ return (self .datapipe , dill_key_value_fn , self ._map , self ._itr , self ._depleted )
143
136
144
137
def __setstate__ (self , state ):
145
- (self .datapipe , dill_key_value_fn , self ._map ) = state
138
+ (self .datapipe , dill_key_value_fn , self ._map , self . _itr , self . _depleted ) = state
146
139
if DILL_AVAILABLE :
147
140
self .key_value_fn = dill .loads (dill_key_value_fn ) # type: ignore[assignment]
148
141
else :
0 commit comments