Skip to content

Commit

Permalink
fix EXHAUSTED check
Browse files Browse the repository at this point in the history
  • Loading branch information
avishaihalev committed Jun 15, 2024
1 parent 1c73be5 commit 9bbeb28
Showing 1 changed file with 8 additions and 3 deletions.
11 changes: 8 additions & 3 deletions mmap_ninja/src/mmap_ninja/parallel.py
Original file line number Diff line number Diff line change
Expand Up @@ -84,7 +84,7 @@ def _collect_no_parallel_batch(self):
results = [_get_from_indexable(self.indexable, j) for j in self._rng()]

if self.exhausted(results):
results = [r for r in results if r != EXHAUSTED]
results = [r for r in results if not isinstance(r, str) or r != EXHAUSTED]

return results

Expand All @@ -94,13 +94,18 @@ def _collect_parallel_batch(self):
results = self._parallel(func(j) for j in self._rng())

if self.exhausted(results):
results = [r for r in results if r != EXHAUSTED]
results = [r for r in results if not isinstance(r, str) or r != EXHAUSTED]
self._parallel.__exit__(None, None, None)

return results

def exhausted(self, results=()):
self._exhausted = self._exhausted or any(r == EXHAUSTED for r in results) or self.completed_batches()
self._exhausted = (
self._exhausted or
any(isinstance(r, str) and r == EXHAUSTED for r in results) or
self.completed_batches()
)

return self._exhausted

def completed_batches(self):
Expand Down

0 comments on commit 9bbeb28

Please sign in to comment.