1"""Non-pipable creation operators."""

2 from __future__ import annotations

3

4 import sys

5 import asyncio

6 import inspect

7 import builtins

8 import itertools

9

10 from typing import (

11 AsyncIterable,

12 Awaitable,

13 Iterable,

14 Protocol,

15 TypeVar,

16 AsyncIterator,

17 cast,

18)

19 from typing_extensions import ParamSpec

20

21 from ..stream import time

22 from ..core import operator, streamcontext

23

24__all__ = [

25 "iterate",

26 "preserve",

27 "just",

28 "call",

29 "throw",

30 "empty",

31 "never",

32 "repeat",

33 "range",

34 "count",

35]

36

37 T = TypeVar("T")

38 P = ParamSpec("P")

39

40# Hack for python 3.8 compatibility

41 if sys.version_info < (3, 9):

42 P = TypeVar("P")

43

44# Convert regular iterables

45

46

47@operator

48 async def from_iterable(it: Iterable[T]) -> AsyncIterator[T]:

49 """Generate values from a regular iterable."""

50 for item in it:

51 await asyncio.sleep(0)

52 yield item

53

54

55@operator

56 def from_async_iterable(ait: AsyncIterable[T]) -> AsyncIterator[T]:

57 """Generate values from an asynchronous iterable.

58

59 Note: the corresponding iterator will be explicitely closed

60 when leaving the context manager."""

61 return streamcontext(ait)

62

63

64@operator

65 def iterate(it: AsyncIterable[T] | Iterable[T]) -> AsyncIterator[T]:

66 """Generate values from a sychronous or asynchronous iterable."""

67 if isinstance(it, AsyncIterable):

68 return from_async_iterable.raw(it)

69 if isinstance(it, Iterable):

70 return from_iterable.raw(it)

71 raise TypeError(f"{type(it).__name__!r} object is not (async) iterable")

72

73

74@operator

75 async def preserve(ait: AsyncIterable[T]) -> AsyncIterator[T]:

76 """Generate values from an asynchronous iterable without

77 explicitly closing the corresponding iterator."""

78 async for item in ait:

79 yield item

80

81

82# Simple operators

83

84

85@operator

86 async def just(value: T) -> AsyncIterator[T]:

87 """Await if possible, and generate a single value."""

88 if inspect.isawaitable(value):

89 yield await value

90 else:

91 yield value

92

93

94 Y = TypeVar("Y", covariant=True)

95

96

97 class SyncCallable(Protocol[P, Y]):

98 def __call__(self, *args: P.args, **kwargs: P.kwargs) -> Y:

99 ...

100

101

102 class AsyncCallable(Protocol[P, Y]):

103 def __call__(self, *args: P.args, **kwargs: P.kwargs) -> Awaitable[Y]:

104 ...

105

106

107@operator

108 async def call(

109 func: SyncCallable[P, T] | AsyncCallable[P, T], *args: P.args, **kwargs: P.kwargs

110) -> AsyncIterator[T]:

111 """Call the given function and generate a single value.

112

113 Await if the provided function is asynchronous.

114 """

115 if asyncio.iscoroutinefunction(func):

116 async_func = cast("AsyncCallable[P, T]", func)

117 yield await async_func(*args, **kwargs)

118 else:

119 sync_func = cast("SyncCallable[P, T]", func)

120 yield sync_func(*args, **kwargs)

121

122

123@operator

124 async def throw(exc: Exception) -> AsyncIterator[None]:

125 """Throw an exception without generating any value."""

126 if False:

127 yield

128 raise exc

129

130

131@operator

132 async def empty() -> AsyncIterator[None]:

133 """Terminate without generating any value."""

134 if False:

135 yield

136

137

138@operator

139 async def never() -> AsyncIterator[None]:

140 """Hang forever without generating any value."""

141 if False:

142 yield

143 future: asyncio.Future[None] = asyncio.Future()

144 try:

145 await future

146 finally:

147 future.cancel()

148

149

150@operator

151 def repeat(

152 value: T, times: int | None = None, *, interval: float = 0.0

153) -> AsyncIterator[T]:

154 """Generate the same value a given number of times.

155

156 If ``times`` is ``None``, the value is repeated indefinitely.

157 An optional interval can be given to space the values out.

158 """

159 args = () if times is None else (times,)

160 it = itertools.repeat(value, *args)

161 agen = from_iterable.raw(it)

162 return time.spaceout.raw(agen, interval) if interval else agen

163

164

165# Counting operators

166

167

168@operator

169 def range(*args: int, interval: float = 0.0) -> AsyncIterator[int]:

170 """Generate a given range of numbers.

171

172 It supports the same arguments as the builtin function.

173 An optional interval can be given to space the values out.

174 """

175 agen = from_iterable.raw(builtins.range(*args))

176 return time.spaceout.raw(agen, interval) if interval else agen

177

178

179@operator

180 def count(

181 start: int = 0, step: int = 1, *, interval: float = 0.0

182) -> AsyncIterator[int]:

183 """Generate consecutive numbers indefinitely.

184

185 Optional starting point and increment can be defined,

186 respectively defaulting to ``0`` and ``1``.

187

188 An optional interval can be given to space the values out.

189 """

190 agen = from_iterable.raw(itertools.count(start, step))

191 return time.spaceout.raw(agen, interval) if interval else agen