"""
 Copyright 2018 Johns Hopkins University  (Author: Jesus Villalba)
 Apache 2.0  (http://www.apache.org/licenses/LICENSE-2.0)
"""

import threading

class ThreadSafeIter:
    """
    Takes an iterator/generator and makes it thread-safe by
    serializing call to the `next` method of given iterator/generator.
    """
    def __init__(self, it):
        self.it = it
        self.lock = threading.Lock()
        
    def __iter__(self):
        return self

    def __next__(self):
        with self.lock:
            return self.it.__next__()
    
    def next(self):
        with self.lock:
            return self.it.next()


def threadsafe_generator(f):
    """ 
    A decorator that takes a generator function and makes it thread-safe.
    """

    def generator(*args, **kwargs):
        return ThreadSafeIter(f(*args, **kwargs))

    return generator

