make generator safe.

generator_safety_2.py

"""                                                                                                                                                                                                                                                                        
generator are not thread safe, two or more thread accessing the same generator                                                                                                                                                                                             
causes an error:                                                                                                                                                                                                                                                           
                                                                                                                                                                                                                                                                           
    ValueError: generator already executing                                                                                                                                                                                                                                
                                                                                                                                                                                                                                                                           
example comes from:                                                                                                                                                                                                                                                        
http://stackoverflow.com/questions/20043472/python-multithreading-why-generators-are-not-thread-safe-what-happens-when-it                                                                                                                                                  
                                                                                                                                                                                                                                                                           
                                                                                                                                                                                                                                                                           
but you can make it safe with a lock                                                                                                                                                                                                                                       
the idea comes while reading joblib source code                                                                                                                                                                                                                            
                                                                                                                                                                                                                                                                           
"""                                                                                                                                                                                                                                                                        
import threading                                                                                                                                                                                                                                                           
                                                                                                                                                                                                                                                                           
                                                                                                                                                                                                                                                                           
class LockedIterator(object):                                                                                                                                                                                                                                              
    def __init__(self, it):                                                                                                                                                                                                                                                
        self._lock = threading.Lock()                                                                                                                                                                                                                                      
        self._it = it                                                                                                                                                                                                                                                      
                                                                                                                                                                                                                                                                           
    def __iter__(self):                                                                                                                                                                                                                                                    
        return self                                                                                                                                                                                                                                                        
                                                                                                                                                                                                                                                                           
    def next(self):                                                                                                                                                                                                                                                        
        with self._lock:                                                                                                                                                                                                                                                   
            return next(self._it)                                                                                                                                                                                                                                          
                                                                                                                                                                                                                                                                           
    # python 3 compat                                                                                                                                                                                                                                                      
    __next__ = next                                                                                                                                                                                                                                                        
                                                                                                                                                                                                                                                                           
                                                                                                                                                                                                                                                                           
class CountThread(threading.Thread):                                                                                                                                                                                                                                       
    def __init__(self, gen):                                                                                                                                                                                                                                               
        super(CountThread, self).__init__()                                                                                                                                                                                                                                
        self.gen = gen                                                                                                                                                                                                                                                     
        self.number_seen = 0                                                                                                                                                                                                                                               
                                                                                                                                                                                                                                                                           
    def run(self):                                                                                                                                                                                                                                                         
        # does nothing but counting                                                                                                                                                                                                                                        
        for i in self.gen:                                                                                                                                                                                                                                                 
            self.number_seen += 1                                                                                                                                                                                                                                          
                                                                                                                                                                                                                                                                           
    def __repr__(self):                                                                                                                                                                                                                                                    
        return "(seen: {})".format(self.number_seen)                                                                                                                                                                                                                       
                                                                                                                                                                                                                                                                           
                                                                                                                                                                                                                                                                           
# shortcut generator expression, or you can write with yeild                                                                                                                                                                                                               
igen = (i for i in xrange(10000))                                                                                                                                                                                                                                          
igen_safe = LockedIterator(igen)                                                                                                                                                                                                                                           
t = [CountThread(igen_safe), CountThread(igen_safe)]                                                                                                                                                                                                                       
                                                                                                                                                                                                                                                                           
[tt.start() for tt in t]                                                                                                                                                                                                                                                   
[tt.join() for tt in t]                                                                                                                                                                                                                                                    
                                                                                                                                                                                                                                                                           
print 'sum should be 10000'                                                                                                                                                                                                                                                
print t