functional.py 2.1 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475
  1. import collections
  2. import functools
  3. import itertools
  4. from typing import ( # noqa: F401
  5. Any,
  6. Callable,
  7. Dict,
  8. Iterable,
  9. List,
  10. Mapping,
  11. Set,
  12. Tuple,
  13. TypeVar,
  14. Union,
  15. )
  16. from .toolz import (
  17. compose as _compose,
  18. )
  19. T = TypeVar("T")
  20. def identity(value: T) -> T:
  21. return value
  22. TGIn = TypeVar("TGIn")
  23. TGOut = TypeVar("TGOut")
  24. TFOut = TypeVar("TFOut")
  25. def combine(
  26. f: Callable[[TGOut], TFOut], g: Callable[[TGIn], TGOut]
  27. ) -> Callable[[TGIn], TFOut]:
  28. return lambda x: f(g(x))
  29. def apply_to_return_value(
  30. callback: Callable[..., T]
  31. ) -> Callable[..., Callable[..., T]]:
  32. def outer(fn: Callable[..., T]) -> Callable[..., T]:
  33. # We would need to type annotate *args and **kwargs but doing so segfaults
  34. # the PyPy builds. We ignore instead.
  35. @functools.wraps(fn)
  36. def inner(*args, **kwargs) -> T: # type: ignore
  37. return callback(fn(*args, **kwargs))
  38. return inner
  39. return outer
  40. TVal = TypeVar("TVal")
  41. TKey = TypeVar("TKey")
  42. to_tuple = apply_to_return_value(
  43. tuple
  44. ) # type: Callable[[Callable[..., Iterable[TVal]]], Callable[..., Tuple[TVal, ...]]] # noqa: E501
  45. to_list = apply_to_return_value(
  46. list
  47. ) # type: Callable[[Callable[..., Iterable[TVal]]], Callable[..., List[TVal]]] # noqa: E501
  48. to_set = apply_to_return_value(
  49. set
  50. ) # type: Callable[[Callable[..., Iterable[TVal]]], Callable[..., Set[TVal]]] # noqa: E501
  51. to_dict = apply_to_return_value(
  52. dict
  53. ) # type: Callable[[Callable[..., Iterable[Union[Mapping[TKey, TVal], Tuple[TKey, TVal]]]]], Callable[..., Dict[TKey, TVal]]] # noqa: E501
  54. to_ordered_dict = apply_to_return_value(
  55. collections.OrderedDict
  56. ) # type: Callable[[Callable[..., Iterable[Union[Mapping[TKey, TVal], Tuple[TKey, TVal]]]]], Callable[..., collections.OrderedDict[TKey, TVal]]] # noqa: E501
  57. sort_return = _compose(to_tuple, apply_to_return_value(sorted))
  58. flatten_return = _compose(
  59. to_tuple, apply_to_return_value(itertools.chain.from_iterable)
  60. )
  61. reversed_return = _compose(to_tuple, apply_to_return_value(reversed), to_tuple)