Only round subcounters when bar is full

Rockhopper-Technologies · Sep 6, 2024 · 11044ce · 11044ce
1 parent 9db944c
commit 11044ce
Show file tree

Hide file tree

Showing 2 changed files with 44 additions and 10 deletions.
diff --git a/enlighten/_counter.py b/enlighten/_counter.py
@@ -733,10 +733,10 @@ def _format_bar(self, fields, iterations, width, elapsed, force_float):
                 remaining.append((remainder, idx))
 
             # Until blocks are accounted for, add full blocks for highest remainders
-            remaining.sort()
-            while sum(block_count) < barLen and remaining:
-                remainder, idx = remaining.pop()
-                if remainder >= 0.85:
+            if self._count == self.total:
+                remaining.sort()
+                while sum(block_count) < barLen and remaining:
+                    _, idx = remaining.pop()
                     block_count[idx] += 1
 
             # Format partial bars

diff --git a/tests/test_subcounter.py b/tests/test_subcounter.py
@@ -358,20 +358,54 @@ def test_subcounter_count_0(self):
         bartext = term.red(BLOCK) + term.blue(BLOCK*3) + term.yellow(BLOCK*35) + ' ' * 41
         self.assertEqual(formatted, bartext)
 
-    def test_subcounter_roundinf(self):
+    def test_subcounter_rounding(self):
         """
-        Extend subcounters to account for remainders
+        Extend subcounters to account for remainders when count reaches total
         """
 
         ctr = self.manager.counter(stream=self.tty.stdout, total=300, bar_format=u'{bar}')
         term = ctr.manager.term
         ctr.count = 151
-        ctr.add_subcounter('yellow', count=132)
-        ctr.add_subcounter('blue', count=12)
-        ctr.add_subcounter('red', count=7)
+        sub1 = ctr.add_subcounter('yellow', count=132)
+        sub2 = ctr.add_subcounter('blue', count=12)
+        sub3 = ctr.add_subcounter('red', count=7)
 
         formatted = ctr.format(width=80)
-        bartext = term.red(BLOCK*2) + term.blue(BLOCK*3) + term.yellow(BLOCK*35) + ' ' * 40
+        bartext = term.red(BLOCK) + term.blue(BLOCK * 3) + term.yellow(BLOCK * 35) + ' ' * 41
+        self.assertEqual(formatted, bartext)
+
+        # Bar complete
+        ctr.count = 300
+        sub1.count = 262
+        sub2.count = 24
+        sub3.count = 14
+
+        formatted = ctr.format(width=80)
+        bartext = term.red(BLOCK * 4) + term.blue(BLOCK * 6) + term.yellow(BLOCK * 70)
+        self.assertEqual(formatted, bartext)
+
+    def test_subcounter_rounding_with_main(self):
+        """
+        Extend subcounters and main counter to account for remainders when count reaches total
+        """
+
+        ctr = self.manager.counter(stream=self.tty.stdout, total=300, bar_format=u'{bar}')
+        term = ctr.manager.term
+        ctr.count = 151
+        sub1 = ctr.add_subcounter('yellow', count=132)
+        sub2 = ctr.add_subcounter('blue', count=12)
+
+        formatted = ctr.format(width=80)
+        bartext = term.blue(BLOCK * 3) + term.yellow(BLOCK * 35) + BLOCK + ' ' * 41
+        self.assertEqual(formatted, bartext)
+
+        # Bar complete
+        ctr.count = 300
+        sub1.count = 262
+        sub2.count = 24
+
+        formatted = ctr.format(width=80)
+        bartext = term.blue(BLOCK * 6) + term.yellow(BLOCK * 70) + BLOCK * 4
         self.assertEqual(formatted, bartext)
 
     def test_subcounter_prefixed(self):