Guard qsize calls

mobiusklein · May 22, 2024 · 3962100 · 3962100
1 parent 691994d
commit 3962100
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 3 deletions.
diff --git a/src/ms_deisotope/tools/deisotoper/collator.py b/src/ms_deisotope/tools/deisotoper/collator.py
@@ -300,11 +300,12 @@ def __iter__(self) -> Iterator[T]:
             if self.consume():
                 self.count_jobs_done += 1
                 try:
-                    if self.queue.qsize() > 500:
+                    if self.queue.qsize() > 500 and self.last_index is not None:
                         self.drain_queue()
                 except NotImplementedError:
                     # Some platforms do not support qsize
-                    self.drain_queue()
+                    if self.last_index is not None:
+                        self.drain_queue()
             if self.last_index is None:
                 keys = sorted(self.waiting)
                 if keys:

diff --git a/src/ms_deisotope/tools/deisotoper/output.py b/src/ms_deisotope/tools/deisotoper/output.py
@@ -135,7 +135,12 @@ def drain_queue():
                     has_work = False
                     continue
                 self._save_bunch(*next_bunch)
-                if self.queue.qsize() > 0:
+                do_drain_queue = True
+                try:
+                    do_drain_queue = self.queue.qsize() > 0
+                except NotImplementedError:
+                    pass
+                if do_drain_queue:
                     current_work = drain_queue()
                     for next_bunch in current_work:
                         i += 1
@@ -144,6 +149,7 @@ def drain_queue():
                         else:
                             self._save_bunch(*next_bunch)
                             i += 1
+
             except QueueEmptyException:
                 continue
             except KeyboardInterrupt: