@@ -1184,33 +1184,28 @@ mod decl {
1184
1184
#[ derive( Debug ) ]
1185
1185
struct PyItertoolsTeeData {
1186
1186
iterable : PyIter ,
1187
- values : PyRwLock < Vec < PyObjectRef > > ,
1188
- locked : AtomicCell < bool > ,
1187
+ values : PyMutex < Vec < PyObjectRef > > ,
1189
1188
}
1190
1189
1191
1190
impl PyItertoolsTeeData {
1192
1191
fn new ( iterable : PyIter , _vm : & VirtualMachine ) -> PyResult < PyRc < Self > > {
1193
1192
Ok ( PyRc :: new ( Self {
1194
1193
iterable,
1195
- values : PyRwLock :: new ( vec ! [ ] ) ,
1196
- locked : AtomicCell :: new ( false ) ,
1194
+ values : PyMutex :: new ( vec ! [ ] ) ,
1197
1195
} ) )
1198
1196
}
1199
1197
1200
1198
fn get_item ( & self , vm : & VirtualMachine , index : usize ) -> PyResult < PyIterReturn > {
1201
- if self . values . read ( ) . len ( ) == index {
1202
- if self . locked . swap ( true ) {
1203
- return Err ( vm. new_runtime_error ( "cannot re-enter the tee iterator" ) ) ;
1204
- }
1205
-
1206
- let result = self . iterable . next ( vm) ;
1207
- self . locked . store ( false ) ;
1199
+ let Some ( mut values) = self . values . try_lock ( ) else {
1200
+ return Err ( vm. new_runtime_error ( "cannot re-enter the tee iterator" ) ) ;
1201
+ } ;
1208
1202
1209
- let obj = raise_if_stop ! ( result?) ;
1210
- self . values . write ( ) . push ( obj) ;
1203
+ if values. len ( ) == index {
1204
+ let obj = raise_if_stop ! ( self . iterable. next( vm) ?) ;
1205
+ values. push ( obj) ;
1211
1206
}
1212
1207
1213
- Ok ( PyIterReturn :: Return ( self . values . read ( ) [ index] . clone ( ) ) )
1208
+ Ok ( PyIterReturn :: Return ( values[ index] . clone ( ) ) )
1214
1209
}
1215
1210
}
1216
1211
0 commit comments