|
17 | 17 |
|
18 | 18 |
|
19 | 19 | from collections import deque |
| 20 | +from warnings import warn |
20 | 21 |
|
21 | 22 | from ..._async_compat.util import AsyncUtil |
22 | 23 | from ...data import DataDehydrator |
@@ -248,11 +249,11 @@ async def _buffer(self, n=None): |
248 | 249 | record_buffer.append(record) |
249 | 250 | if n is not None and len(record_buffer) >= n: |
250 | 251 | break |
251 | | - self._exhausted = False |
252 | 252 | if n is None: |
253 | 253 | self._record_buffer = record_buffer |
254 | 254 | else: |
255 | 255 | self._record_buffer.extend(record_buffer) |
| 256 | + self._exhausted = not self._record_buffer |
256 | 257 |
|
257 | 258 | async def _buffer_all(self): |
258 | 259 | """Sets the Result object in an detached state by fetching all records |
@@ -294,6 +295,14 @@ async def _tx_end(self): |
294 | 295 | await self.consume() |
295 | 296 | self._out_of_scope = True |
296 | 297 |
|
| 298 | + async def _exhaust(self): |
| 299 | + # Exhaust the result, ditching all remaining records. |
| 300 | + if not self._exhausted: |
| 301 | + self._discarding = True |
| 302 | + self._record_buffer.clear() |
| 303 | + async for _ in self: |
| 304 | + pass |
| 305 | + |
297 | 306 | async def consume(self): |
298 | 307 | """Consume the remainder of this result and return a :class:`neo4j.ResultSummary`. |
299 | 308 |
|
@@ -330,42 +339,85 @@ async def get_two_tx(tx): |
330 | 339 |
|
331 | 340 | :returns: The :class:`neo4j.ResultSummary` for this result |
332 | 341 | """ |
333 | | - if self._exhausted is False: |
334 | | - self._discarding = True |
335 | | - async for _ in self: |
336 | | - pass |
| 342 | + if self._exhausted: |
| 343 | + if self._out_of_scope: |
| 344 | + raise ResultConsumedError(self, _RESULT_OUT_OF_SCOPE_ERROR) |
| 345 | + if self._consumed: |
| 346 | + raise ResultConsumedError(self, _RESULT_CONSUMED_ERROR) |
| 347 | + else: |
| 348 | + await self._exhaust() |
337 | 349 |
|
338 | 350 | summary = self._obtain_summary() |
339 | 351 | self._consumed = True |
340 | 352 | return summary |
341 | 353 |
|
342 | | - async def single(self): |
343 | | - """Obtain the next and only remaining record from this result if available else return None. |
| 354 | + async def single(self, strict=False): |
| 355 | + """Obtain the next and only remaining record or None. |
| 356 | +
|
344 | 357 | Calling this method always exhausts the result. |
345 | 358 |
|
346 | 359 | A warning is generated if more than one record is available but |
347 | 360 | the first of these is still returned. |
348 | 361 |
|
349 | | - :returns: the next :class:`neo4j.AsyncRecord`. |
| 362 | + :param strict: |
| 363 | + If :const:`True`, raise a :class:`neo4j.ResultNotSingleError` |
| 364 | + instead of returning None if there is more than one record or |
| 365 | + warning if there are more than 1 record. |
| 366 | + :const:`False` by default. |
| 367 | + :type strict: bool |
350 | 368 |
|
351 | | - :raises ResultNotSingleError: if not exactly one record is available. |
352 | | - :raises ResultConsumedError: if the transaction from which this result was |
353 | | - obtained has been closed. |
| 369 | + :returns: the next :class:`neo4j.Record` or :const:`None` if none remain |
| 370 | + :warns: if more than one record is available |
| 371 | +
|
| 372 | + :raises ResultNotSingleError: |
| 373 | + If ``strict=True`` and not exactly one record is available. |
| 374 | + :raises ResultConsumedError: if the transaction from which this result |
| 375 | + was obtained has been closed. |
| 376 | +
|
| 377 | + .. versionchanged:: 5.0 |
| 378 | + Added ``strict`` parameter. |
354 | 379 | """ |
355 | 380 | await self._buffer(2) |
356 | | - if not self._record_buffer: |
| 381 | + buffer = self._record_buffer |
| 382 | + self._record_buffer = deque() |
| 383 | + await self._exhaust() |
| 384 | + if not buffer: |
| 385 | + if not strict: |
| 386 | + return None |
357 | 387 | raise ResultNotSingleError( |
358 | 388 | self, |
359 | 389 | "No records found. " |
360 | 390 | "Make sure your query returns exactly one record." |
361 | 391 | ) |
362 | | - elif len(self._record_buffer) > 1: |
363 | | - raise ResultNotSingleError( |
364 | | - self, |
365 | | - "More than one record found. " |
366 | | - "Make sure your query returns exactly one record." |
367 | | - ) |
368 | | - return self._record_buffer.popleft() |
| 392 | + elif len(buffer) > 1: |
| 393 | + res = buffer.popleft() |
| 394 | + if not strict: |
| 395 | + warn("Expected a result with a single record, " |
| 396 | + "but found multiple.") |
| 397 | + return res |
| 398 | + else: |
| 399 | + raise ResultNotSingleError( |
| 400 | + self, |
| 401 | + "More than one record found. " |
| 402 | + "Make sure your query returns exactly one record." |
| 403 | + ) |
| 404 | + return buffer.popleft() |
| 405 | + |
| 406 | + async def fetch(self, n): |
| 407 | + """Obtain up to n records from this result. |
| 408 | +
|
| 409 | + :param n: the maximum number of records to fetch. |
| 410 | + :type n: int |
| 411 | +
|
| 412 | + :returns: list of :class:`neo4j.AsyncRecord` |
| 413 | +
|
| 414 | + .. versionadded:: 5.0 |
| 415 | + """ |
| 416 | + await self._buffer(n) |
| 417 | + return [ |
| 418 | + self._record_buffer.popleft() |
| 419 | + for _ in range(min(n, len(self._record_buffer))) |
| 420 | + ] |
369 | 421 |
|
370 | 422 | async def peek(self): |
371 | 423 | """Obtain the next record from this result without consuming it. |
|
0 commit comments