Ensure streams are flushed consistantly

Rockhopper-Technologies · Aug 12, 2020 · ccf77ff · ccf77ff
1 parent 9de756f
commit ccf77ff
Show file tree

Hide file tree

Showing 2 changed files with 19 additions and 12 deletions.
diff --git a/enlighten/_manager.py b/enlighten/_manager.py
@@ -250,8 +250,8 @@ def _add_counter(self, counter_class, *args, **kwargs):
 
         self._set_scroll_area()
         for ctr in reversed(toRefresh):
-            ctr.refresh()
-        self.stream.flush()
+            ctr.refresh(flush=False)
+        self._flush_streams()
 
         return new
 
@@ -293,7 +293,7 @@ def _resize_handler(self, *args, **kwarg):  # pylint: disable=unused-argument
 
             for counter in self.counters:
                 counter.refresh(flush=False)
-            self.stream.flush()
+            self._flush_streams()
 
             self.resize_lock = False
 
@@ -346,6 +346,15 @@ def _set_scroll_area(self, force=False):
             if self.companion_term is not None:
                 self.companion_term.move_to(0, scrollPosition)
 
+    def _flush_streams(self):
+        """
+        Convenience method for flushing streams
+        """
+
+        self.stream.flush()
+        if self.companion_stream is not None:
+            self.companion_stream.flush()
+
     def _at_exit(self):
         """
         Resets terminal to normal configuration
@@ -364,9 +373,7 @@ def _at_exit(self):
 
             self.term.feed()
 
-            self.stream.flush()
-            if self.companion_stream is not None:
-                self.companion_stream.flush()
+            self._flush_streams()
 
         except ValueError:  # Possibly closed file handles
             pass
@@ -420,8 +427,6 @@ def stop(self):
                     term.move_to(0, term.height - num)
                     stream.write(term.clear_eol)
 
-            stream.flush()
-
         finally:
 
             if self.set_scroll:
@@ -443,6 +448,8 @@ def stop(self):
         if 1 in positions:
             term.feed()
 
+        self._flush_streams()
+
     def write(self, output='', flush=True, counter=None):
         """
         Args:
@@ -472,7 +479,7 @@ def write(self, output='', flush=True, counter=None):
                     self._autorefresh(exclude=(counter,))
                 self._set_scroll_area()
                 if flush:
-                    stream.flush()
+                    self._flush_streams()
 
     def _autorefresh(self, exclude):
         """

diff --git a/tests/test_manager.py b/tests/test_manager.py
@@ -136,7 +136,7 @@ def test_counter_and_remove(self):
         self.assertEqual(manager.counters[counter1], 2)
         self.assertEqual(manager.counters[counter2], 1)
         self.assertEqual(counter1.calls,
-                         ['clear(flush=False)', 'refresh(flush=True, elapsed=None)'])
+                         ['clear(flush=False)', 'refresh(flush=False, elapsed=None)'])
         self.assertEqual(counter2.calls, [])
         self.assertEqual(ssa.call_count, 1)
         counter1.calls = []
@@ -149,9 +149,9 @@ def test_counter_and_remove(self):
         self.assertEqual(manager.counters[counter2], 2)
         self.assertEqual(manager.counters[counter3], 1)
         self.assertEqual(counter1.calls,
-                         ['clear(flush=False)', 'refresh(flush=True, elapsed=None)'])
+                         ['clear(flush=False)', 'refresh(flush=False, elapsed=None)'])
         self.assertEqual(counter2.calls,
-                         ['clear(flush=False)', 'refresh(flush=True, elapsed=None)'])
+                         ['clear(flush=False)', 'refresh(flush=False, elapsed=None)'])
         self.assertEqual(counter3.calls, [])
         self.assertEqual(ssa.call_count, 1)
         counter1.calls = []