diff --git a/pims/pyav_reader.py b/pims/pyav_reader.py index 2b453ee..40caa1e 100644 --- a/pims/pyav_reader.py +++ b/pims/pyav_reader.py @@ -2,6 +2,7 @@ unicode_literals) import numpy as np +import threading from pims.base_frames import FramesSequence from pims.frame import Frame @@ -27,6 +28,28 @@ def _next_video_packet(container_iter): raise ValueError("Could not find any video packets.") +class threadsafe_iter: + """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 next(self.it) + +def threadsafe_generator(f): + """A decorator that takes a generator function and makes it thread-safe. + """ + def g(*a, **kw): + return threadsafe_iter(f(*a, **kw)) + return g + class WrapPyAvFrame(object): def __init__(self, frame, frame_no, metadata=None): self.frame_no = frame_no @@ -46,7 +69,7 @@ def to_frame(self): frame_no=self.frame_no, metadata=self.metadata) return self.arr - +@threadsafe_generator def _gen_frames(demuxer, time_base, frame_rate=1., first_pts=0): for packet in demuxer: for frame in packet.decode():