utils_worker.py 1.2 KB

1234567891011121314151617181920212223242526272829303132333435363738
  1. """
  2. IO/concurrency helpers for `tqdm.contrib`.
  3. """
  4. from collections import deque
  5. from concurrent.futures import ThreadPoolExecutor
  6. from ..auto import tqdm as tqdm_auto
  7. __author__ = {"github.com/": ["casperdcl"]}
  8. __all__ = ['MonoWorker']
  9. class MonoWorker(object):
  10. """
  11. Supports one running task and one waiting task.
  12. The waiting task is the most recent submitted (others are discarded).
  13. """
  14. def __init__(self):
  15. self.pool = ThreadPoolExecutor(max_workers=1)
  16. self.futures = deque([], 2)
  17. def submit(self, func, *args, **kwargs):
  18. """`func(*args, **kwargs)` may replace currently waiting task."""
  19. futures = self.futures
  20. if len(futures) == futures.maxlen:
  21. running = futures.popleft()
  22. if not running.done():
  23. if len(futures): # clear waiting
  24. waiting = futures.pop()
  25. waiting.cancel()
  26. futures.appendleft(running) # re-insert running
  27. try:
  28. waiting = self.pool.submit(func, *args, **kwargs)
  29. except Exception as e:
  30. tqdm_auto.write(str(e))
  31. else:
  32. futures.append(waiting)
  33. return waiting