123 lines
2.9 KiB
Python
123 lines
2.9 KiB
Python
# Copyright (c) Twisted Matrix Laboratories.
|
|
# See LICENSE for details.
|
|
|
|
"""
|
|
General functional-style helpers for disttrial.
|
|
"""
|
|
|
|
from functools import partial, wraps
|
|
from typing import Awaitable, Callable, Iterable, Optional, TypeVar
|
|
|
|
from twisted.internet.defer import Deferred, succeed
|
|
|
|
_A = TypeVar("_A")
|
|
_B = TypeVar("_B")
|
|
_C = TypeVar("_C")
|
|
|
|
|
|
def fromOptional(default: _A, optional: Optional[_A]) -> _A:
|
|
"""
|
|
Get a definite value from an optional value.
|
|
|
|
@param default: The value to return if the optional value is missing.
|
|
|
|
@param optional: The optional value to return if it exists.
|
|
"""
|
|
if optional is None:
|
|
return default
|
|
return optional
|
|
|
|
|
|
def takeWhile(condition: Callable[[_A], bool], xs: Iterable[_A]) -> Iterable[_A]:
|
|
"""
|
|
:return: An iterable over C{xs} that stops when C{condition} returns
|
|
``False`` based on the value of iterated C{xs}.
|
|
"""
|
|
for x in xs:
|
|
if condition(x):
|
|
yield x
|
|
else:
|
|
break
|
|
|
|
|
|
async def sequence(a: Awaitable[_A], b: Awaitable[_B]) -> _B:
|
|
"""
|
|
Wait for one action to complete and then another.
|
|
|
|
If either action fails, failure is propagated. If the first action fails,
|
|
the second action is not waited on.
|
|
"""
|
|
await a
|
|
return await b
|
|
|
|
|
|
def flip(f: Callable[[_A, _B], _C]) -> Callable[[_B, _A], _C]:
|
|
"""
|
|
Create a function like another but with the order of the first two
|
|
arguments flipped.
|
|
"""
|
|
|
|
@wraps(f)
|
|
def g(b, a):
|
|
return f(a, b)
|
|
|
|
return g
|
|
|
|
|
|
def compose(fx: Callable[[_B], _C], fy: Callable[[_A], _B]) -> Callable[[_A], _C]:
|
|
"""
|
|
Create a function that calls one function with an argument and then
|
|
another function with the result of the first function.
|
|
"""
|
|
|
|
@wraps(fx)
|
|
@wraps(fy)
|
|
def g(a):
|
|
return fx(fy(a))
|
|
|
|
return g
|
|
|
|
|
|
# Discard the result of an awaitable and substitute None in its place.
|
|
discardResult: Callable[[Awaitable[_A]], Deferred[None]] = compose(
|
|
Deferred.fromCoroutine,
|
|
partial(flip(sequence), succeed(None)),
|
|
)
|
|
|
|
|
|
async def iterateWhile(
|
|
predicate: Callable[[_A], bool],
|
|
action: Callable[[], Awaitable[_A]],
|
|
) -> _A:
|
|
"""
|
|
Call a function repeatedly until its result fails to satisfy a predicate.
|
|
|
|
@param predicate: The check to apply.
|
|
|
|
@param action: The function to call.
|
|
|
|
@return: The result of C{action} which did not satisfy C{predicate}.
|
|
"""
|
|
while True:
|
|
result = await action()
|
|
if not predicate(result):
|
|
return result
|
|
|
|
|
|
def countingCalls(f: Callable[[int], _A]) -> Callable[[], _A]:
|
|
"""
|
|
Wrap a function with another that automatically passes an integer counter
|
|
of the number of calls that have gone through the wrapper.
|
|
"""
|
|
counter = 0
|
|
|
|
def g() -> _A:
|
|
nonlocal counter
|
|
try:
|
|
result = f(counter)
|
|
finally:
|
|
counter += 1
|
|
return result
|
|
|
|
return g
|