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