smart_future.py

   1 #!/usr/bin/env python3
   2
   3 from __future__ import annotations
   4 from collections.abc import Mapping
   5 import concurrent.futures as fut
   6 import time
   7 from typing import Callable, List, TypeVar
   8
   9 # This module is commonly used by others in here and should avoid
  10 # taking any unnecessary dependencies back on them.
  11 from deferred_operand import DeferredOperand
  12 import id_generator
  13
  14 T = TypeVar('T')
  15
  16
  17 def wait_any(futures: List[SmartFuture], *, callback: Callable = None):
  18     finished: Mapping[int, bool] = {}
  19     x = 0
  20     while True:
  21         future = futures[x]
  22         if not finished.get(future.get_id(), False):
  23             if future.is_ready():
  24                 finished[future.get_id()] = True
  25                 yield future
  26             else:
  27                 if callback is not None:
  28                     callback()
  29                 time.sleep(0.1)
  30         x += 1
  31         if x >= len(futures):
  32             x = 0
  33         if len(finished) == len(futures):
  34             if callback is not None:
  35                 callback()
  36             return
  37
  38
  39 def wait_all(futures: List[SmartFuture]) -> None:
  40     done_set = set()
  41     while len(done_set) < len(futures):
  42         for future in futures:
  43             i = future.get_id()
  44             if i not in done_set and future.wrapped_future.done():
  45                 done_set.add(i)
  46             time.sleep(0.1)
  47
  48
  49 class SmartFuture(DeferredOperand):
  50     """This is a SmartFuture, a class that wraps a normal Future and can
  51     then be used, mostly, like a normal (non-Future) identifier.
  52
  53     Using a FutureWrapper in expressions will block and wait until
  54     the result of the deferred operation is known.
  55     """
  56
  57     def __init__(self, wrapped_future: fut.Future) -> None:
  58         self.wrapped_future = wrapped_future
  59         self.id = id_generator.get("smart_future_id")
  60
  61     def get_id(self) -> int:
  62         return self.id
  63
  64     def is_ready(self) -> bool:
  65         return self.wrapped_future.done()
  66
  67     # You shouldn't have to call this; instead, have a look at defining a
  68     # method on DeferredOperand base class.
  69     def _resolve(self, *, timeout=None) -> T:
  70         return self.wrapped_future.result(timeout)