Coverage for src / lexigram / contracts / infra / streams.py: 28%

1"""Typed async stream primitives.""" 

2 

3from __future__ import annotations 

4 

5from collections.abc import AsyncIterator, Awaitable, Callable 

6from typing import Any, Generic, TypeVar, cast 

7 

8from lexigram.contracts.core.result import Result 

9 

10T = TypeVar("T") 

11E = TypeVar("E") 

12U = TypeVar("U") 

13F = TypeVar("F") 

14 

15 

16class _StoredStreamError(Exception, Generic[E]): 

17 """Internal carrier for propagating typed stream failures.""" 

18 

19 _code: str = "LEX_ERR_INFRA_011" 

20 error: E 

21 

22 def __init__(self, error: E) -> None: 

23 super().__init__(str(error)) 

24 self.error = error 

25 

26 

27class _StreamOk(Result[T, E]): 

28 """Private ``Ok`` implementation for contracts-only helpers.""" 

29 

30 def __init__(self, value: T) -> None: 

31 self._value = value 

32 

33 def is_ok(self) -> bool: 

34 return True 

35 

36 def is_err(self) -> bool: 

37 return False 

38 

39 def unwrap(self) -> T: 

40 return self._value 

41 

42 def unwrap_err(self) -> E: 

43 raise ValueError("Called unwrap_err() on Ok result") 

44 

45 def unwrap_or(self, default: T) -> T: 

46 return self._value 

47 

48 def unwrap_or_else(self, op: Callable[[E], T]) -> T: 

49 return self._value 

50 

51 def map_sync(self, op: Callable[[T], U]) -> Result[U, E]: 

52 return _StreamOk(op(self._value)) 

53 

54 def map_err(self, op: Callable[[E], F]) -> Result[T, F]: 

55 return cast("Result[T, F]", self) 

56 

57 def and_then_sync(self, op: Callable[[T], Result[U, E]]) -> Result[U, E]: 

58 return op(self._value) 

59 

60 def or_else_sync(self, op: Callable[[E], Result[T, F]]) -> Result[T, F]: 

61 return cast("Result[T, F]", self) 

62 

63 def expect(self, message: str) -> T: 

64 return self._value 

65 

66 def match(self, ok: Callable[[T], U], err: Callable[[E], U]) -> U: 

67 return ok(self._value) 

68 

69 async def map(self, op: Callable[[T], Awaitable[U]]) -> Result[U, E]: 

70 return _StreamOk(await op(self._value)) 

71 

72 async def and_then( 

73 self, op: Callable[[T], Awaitable[Result[U, E]]] 

74 ) -> Result[U, E]: 

75 return await op(self._value) 

76 

77 async def or_else(self, op: Callable[[E], Awaitable[Result[T, F]]]) -> Result[T, F]: 

78 return cast("Result[T, F]", self) 

79 

80 def flatten(self) -> Result[Any, E]: 

81 if isinstance(self._value, Result): 

82 return self._value 

83 return cast("Result[Any, E]", self) 

84 

85 def filter(self, predicate: Callable[[T], bool], error: E) -> Result[T, E]: 

86 return self if predicate(self._value) else _StreamErr(error) 

87 

88 def ok_or(self, default: U) -> T | U: 

89 return self._value 

90 

91 

92class _StreamErr(Result[T, E]): 

93 """Private ``Err`` implementation for contracts-only helpers.""" 

94 

95 def __init__(self, error: E) -> None: 

96 self._error = error 

97 

98 def is_ok(self) -> bool: 

99 return False 

100 

101 def is_err(self) -> bool: 

102 return True 

103 

104 def unwrap(self) -> T: 

105 if isinstance(self._error, Exception): 

106 raise self._error 

107 raise ValueError(str(self._error)) 

108 

109 def unwrap_err(self) -> E: 

110 return self._error 

111 

112 def unwrap_or(self, default: T) -> T: 

113 return default 

114 

115 def unwrap_or_else(self, op: Callable[[E], T]) -> T: 

116 return op(self._error) 

117 

118 def map_sync(self, op: Callable[[T], U]) -> Result[U, E]: 

119 return cast("Result[U, E]", self) 

120 

121 def map_err(self, op: Callable[[E], F]) -> Result[T, F]: 

122 return _StreamErr(op(self._error)) 

123 

124 def and_then_sync(self, op: Callable[[T], Result[U, E]]) -> Result[U, E]: 

125 return cast("Result[U, E]", self) 

126 

127 def or_else_sync(self, op: Callable[[E], Result[T, F]]) -> Result[T, F]: 

128 return op(self._error) 

129 

130 def expect(self, message: str) -> T: 

131 raise ValueError(message) 

132 

133 def match(self, ok: Callable[[T], U], err: Callable[[E], U]) -> U: 

134 return err(self._error) 

135 

136 async def map(self, op: Callable[[T], Awaitable[U]]) -> Result[U, E]: 

137 return cast("Result[U, E]", self) 

138 

139 async def and_then( 

140 self, op: Callable[[T], Awaitable[Result[U, E]]] 

141 ) -> Result[U, E]: 

142 return cast("Result[U, E]", self) 

143 

144 async def or_else(self, op: Callable[[E], Awaitable[Result[T, F]]]) -> Result[T, F]: 

145 return await op(self._error) 

146 

147 def flatten(self) -> Result[Any, E]: 

148 return cast("Result[Any, E]", self) 

149 

150 def filter(self, predicate: Callable[[T], bool], error: E) -> Result[T, E]: 

151 return self 

152 

153 def ok_or(self, default: U) -> T | U: 

154 return default 

155 

156 

157class AsyncStream(Generic[T, E]): 

158 """Async stream that preserves typed failures for terminal operations. 

159 

160 ``AsyncStream`` is still directly iterable with ``async for``. If the 

161 underlying iterator fails, the stream stores the typed error, stops 

162 iteration, and exposes the failure through ``collect()``, ``first()``, 

163 ``drain()``, and the ``error`` property. 

164 """ 

165 

166 def __init__( 

167 self, 

168 iterator: AsyncIterator[T], 

169 *, 

170 error_adapter: Callable[[Exception], E], 

171 ) -> None: 

172 self._iterator = iterator 

173 self._error_adapter = error_adapter 

174 self._error: E | None = None 

175 self._done = False 

176 

177 @property 

178 def error(self) -> E | None: 

179 """Return the stored stream error, if one occurred.""" 

180 return self._error 

181 

182 @property 

183 def failed(self) -> bool: 

184 """Return ``True`` when stream iteration ended with an error.""" 

185 return self._error is not None 

186 

187 def __aiter__(self) -> AsyncIterator[T]: 

188 return self._iterate() 

189 

190 async def _iterate(self) -> AsyncIterator[T]: 

191 if self._done: 

192 return 

193 

194 while True: 

195 try: 

196 item = await anext(self._iterator) 

197 except StopAsyncIteration: 

198 self._done = True 

199 return 

200 except Exception as exc: # noqa: BLE001 

201 self._done = True 

202 self._error = self._coerce_error(exc) 

203 return 

204 else: 

205 yield item 

206 

207 async def collect(self) -> Result[list[T], E]: 

208 """Consume the entire stream into a list or return the first failure.""" 

209 items: list[T] = [] 

210 async for item in self: 

211 items.append(item) 

212 

213 if self._error is not None: 

214 return _StreamErr(self._error) 

215 return _StreamOk(items) 

216 

217 async def first(self) -> Result[T, E]: 

218 """Return the first item from the stream or the first failure.""" 

219 async for item in self: 

220 return _StreamOk(item) 

221 

222 if self._error is not None: 

223 return _StreamErr(self._error) 

224 raise ValueError("AsyncStream is empty") 

225 

226 def map(self, fn: Callable[[T], U]) -> AsyncStream[U, E]: 

227 """Transform stream items lazily.""" 

228 

229 async def _generator() -> AsyncIterator[U]: 

230 async for item in self: 

231 yield fn(item) 

232 self._raise_if_failed() 

233 

234 return AsyncStream(_generator(), error_adapter=self._coerce_error) 

235 

236 def filter(self, predicate: Callable[[T], bool]) -> AsyncStream[T, E]: 

237 """Filter stream items lazily.""" 

238 

239 async def _generator() -> AsyncIterator[T]: 

240 async for item in self: 

241 if predicate(item): 

242 yield item 

243 self._raise_if_failed() 

244 

245 return AsyncStream(_generator(), error_adapter=self._coerce_error) 

246 

247 def take(self, count: int) -> AsyncStream[T, E]: 

248 """Return a stream containing at most ``count`` items.""" 

249 

250 async def _generator() -> AsyncIterator[T]: 

251 if count <= 0: 

252 return 

253 

254 remaining = count 

255 async for item in self: 

256 yield item 

257 remaining -= 1 

258 if remaining == 0: 

259 return 

260 self._raise_if_failed() 

261 

262 return AsyncStream(_generator(), error_adapter=self._coerce_error) 

263 

264 async def drain(self) -> Result[None, E]: 

265 """Consume the stream without collecting items.""" 

266 async for _ in self: 

267 continue 

268 

269 if self._error is not None: 

270 return _StreamErr(self._error) 

271 return _StreamOk(None) 

272 

273 def _coerce_error(self, exc: Exception) -> E: 

274 if isinstance(exc, _StoredStreamError): 

275 return cast("E", exc.error) 

276 return self._error_adapter(exc) 

277 

278 def _raise_if_failed(self) -> None: 

279 if self._error is not None: 

280 raise _StoredStreamError(self._error) 

281 

282 

283__all__ = ["AsyncStream"]