smart_future.py

   1 #!/usr/bin/env python3
   2
   3 from __future__ import annotations
   4 import concurrent
   5 import concurrent.futures as fut
   6 import logging
   7 import traceback
   8 from typing import Callable, List, TypeVar
   9
  10 from overrides import overrides
  11
  12 # This module is commonly used by others in here and should avoid
  13 # taking any unnecessary dependencies back on them.
  14 from deferred_operand import DeferredOperand
  15 import id_generator
  16
  17 logger = logging.getLogger(__name__)
  18
  19 T = TypeVar('T')
  20
  21
  22 def wait_any(
  23     futures: List[SmartFuture],
  24     *,
  25     callback: Callable = None,
  26     log_exceptions: bool = True,
  27 ):
  28     real_futures = []
  29     smart_future_by_real_future = {}
  30     completed_futures = set()
  31     for _ in futures:
  32         real_futures.append(_.wrapped_future)
  33         smart_future_by_real_future[_.wrapped_future] = _
  34     while len(completed_futures) != len(real_futures):
  35         newly_completed_futures = concurrent.futures.as_completed(real_futures)
  36         for f in newly_completed_futures:
  37             if callback is not None:
  38                 callback()
  39             completed_futures.add(f)
  40             if log_exceptions and not f.cancelled():
  41                 exception = f.exception()
  42                 if exception is not None:
  43                     logger.exception(exception)
  44                     traceback.print_tb(exception.__traceback__)
  45             yield smart_future_by_real_future[f]
  46     if callback is not None:
  47         callback()
  48     return
  49
  50
  51 def wait_all(
  52     futures: List[SmartFuture],
  53     *,
  54     log_exceptions: bool = True,
  55 ) -> None:
  56     real_futures = [x.wrapped_future for x in futures]
  57     (done, not_done) = concurrent.futures.wait(
  58         real_futures, timeout=None, return_when=concurrent.futures.ALL_COMPLETED
  59     )
  60     if log_exceptions:
  61         for f in real_futures:
  62             if not f.cancelled():
  63                 exception = f.exception()
  64                 if exception is not None:
  65                     logger.exception(exception)
  66                     traceback.print_tb(exception.__traceback__)
  67     assert len(done) == len(real_futures)
  68     assert len(not_done) == 0
  69
  70
  71 class SmartFuture(DeferredOperand):
  72     """This is a SmartFuture, a class that wraps a normal Future and can
  73     then be used, mostly, like a normal (non-Future) identifier.
  74
  75     Using a FutureWrapper in expressions will block and wait until
  76     the result of the deferred operation is known.
  77     """
  78
  79     def __init__(self, wrapped_future: fut.Future) -> None:
  80         self.wrapped_future = wrapped_future
  81         self.id = id_generator.get("smart_future_id")
  82
  83     def get_id(self) -> int:
  84         return self.id
  85
  86     def is_ready(self) -> bool:
  87         return self.wrapped_future.done()
  88
  89     # You shouldn't have to call this; instead, have a look at defining a
  90     # method on DeferredOperand base class.
  91     @overrides
  92     def _resolve(self, *, timeout=None) -> T:
  93         return self.wrapped_future.result(timeout)