Repository: Deric-W/ascii_progress

Issue:
improved context managers
Problem:
The messages displayed when using `Spinner` or `Bar` as context managers are hard coded.
 
Solution:
Add a `.handle_exceptions(message, error)` method to both classes which returns a context manager.

Comments:

Patch:
diff --git a/ascii_progress/__main__.py b/ascii_progress/__main__.py
index f5d1e15..b550f91 100644
--- a/ascii_progress/__main__.py
+++ b/ascii_progress/__main__.py
@@ -16,7 +16,7 @@ for frames in (
     (">))'>", "    >))'>", "        >))'>", "        <'((<", "    <'((<", "<'((<")
 ):
     sys.stdout.write("Working ")
-    with Spinner(frames) as spinner:
+    with Spinner(frames).handle_exceptions("Done", "Exception") as spinner:
         for _ in map(spinner.set_progress, range(1, 15)):
             time.sleep(0.2)
 
@@ -29,6 +29,6 @@ for bar_format in map(
     )
 ):
     sys.stdout.write("Working ")
-    with bar_format.bar(75) as bar:
+    with bar_format.bar(75).handle_exceptions("Done", "Exception") as bar:
         for _ in bar:
             time.sleep(0.02)
diff --git a/ascii_progress/bar.py b/ascii_progress/bar.py
index e660116..b5f0903 100644
--- a/ascii_progress/bar.py
+++ b/ascii_progress/bar.py
@@ -9,11 +9,23 @@ import sys
 import math
 from abc import abstractmethod
 from string import Formatter
-from typing import Callable, Sequence, TextIO, Tuple, Iterator, ContextManager, Mapping, Any, Union
+from typing import (
+    Callable,
+    Sequence,
+    TextIO,
+    Tuple,
+    Iterator,
+    ContextManager,
+    Mapping,
+    Any,
+    Union,
+    TypeVar
+)
 
 __all__ = (
     "LazyFormatter",
     "LAZY_FORMATTER",
+    "BarContext",
     "Bar",
     "ThresholdDecorator",
     "PercentDecorator",
@@ -36,8 +48,48 @@ class LazyFormatter(Formatter):
 
 LAZY_FORMATTER = LazyFormatter()
 
+T = TypeVar("T", bound="Bar")
+
+
+class BarContext(ContextManager[T]):
+    """context manager which handles exceptions while using a bar"""
+
+    bar: T
+
+    message: str
+
+    error: str
+
+    __slots__ = ("bar", "message", "error")
+
+    def __init__(self, bar: T, message: str, error: str) -> None:
+        self.bar = bar
+        self.message = message
+        self.error = error
+
+    def __eq__(self, other: object) -> bool:
+        if isinstance(other, BarContext):
+            return self.bar == other.bar \
+                and self.message == other.message \
+                and self.error == other.error
+        return NotImplemented
+
+    def __enter__(self) -> T:
+        return self.bar
+
+    def __exit__(self, type, value, traceback) -> bool:     # type: ignore
+        if type is None:
+            self.bar.replace(self.message, end="\n")
+        elif type is KeyboardInterrupt:
+            # add 2 \b and 2 spaces to handle additional ^C
+            # add 2 additional spaces to make up for missing padding
+            self.bar.replace("\b\b" + self.error, end="    \n")
+        else:
+            self.bar.replace(self.error, end="\n")
+        return False    # we dont handle exceptions
+
 
-class Bar(ContextManager["Bar"], Iterator[None]):
+class Bar(Iterator[None]):
     """abstract base class for progress bars"""
 
     __slots__ = ()
@@ -49,16 +101,6 @@ class Bar(ContextManager["Bar"], Iterator[None]):
                 and self.width() == other.width()
         return NotImplemented
 
-    def __enter__(self) -> "Bar":
-        return self
-
-    def __exit__(self, type, value, traceback) -> bool:     # type: ignore
-        if type is KeyboardInterrupt:   # handle ^C
-            self.replace("\b\bKeyboardInterrupt", end="  \n")
-        else:
-            self.replace("Finished")
-        return False    # we dont handle exceptions
-
     def __iter__(self) -> "Bar":
         return self
 
@@ -97,6 +139,10 @@ class Bar(ContextManager["Bar"], Iterator[None]):
         """return the size of the bar"""
         raise NotImplementedError
 
+    def handle_exceptions(self: T, message: str, error: str) -> BarContext[T]:
+        """return a context manager which replaces the bar with message or error if a exceptions is raised"""
+        return BarContext(self, message, error)
+
     def ratio(self) -> float:
         """return the ration progress / target"""
         return self.progress() / self.target()
@@ -271,16 +317,14 @@ class PercentDecorator(ThresholdDecorator):
     @classmethod
     def with_inferred_thresholds(cls, bar: Bar) -> "PercentDecorator":
         """create an instance with inferred thresholds"""
-        lower_threshold, upper_threshold = calculate_thresholds(bar, 100)
-        return cls(
-            bar,
-            lower_threshold,
-            upper_threshold,
-        )
+        return cls(bar, *calculate_thresholds(bar, 100))
 
     def update_thresholds(self) -> None:
         """update lower and upper thresholds"""
-        self.lower_threshold, self.upper_threshold = calculate_thresholds(self.bar, 100)
+        self.lower_threshold, self.upper_threshold = calculate_thresholds(
+            self.bar,
+            100
+        )
 
 
 class BarDecorator(ThresholdDecorator):
@@ -291,16 +335,14 @@ class BarDecorator(ThresholdDecorator):
     @classmethod
     def with_inferred_thresholds(cls, bar: Bar) -> "BarDecorator":
         """create an instance with inferred thresholds"""
-        lower_threshold, upper_threshold = calculate_thresholds(bar, bar.width())
-        return cls(
-            bar,
-            lower_threshold,
-            upper_threshold,
-        )
+        return cls(bar, *calculate_thresholds(bar, bar.width()))
 
     def update_thresholds(self) -> None:
         """update lower and upper thresholds"""
-        self.lower_threshold, self.upper_threshold = calculate_thresholds(self.bar, self.bar.width())
+        self.lower_threshold, self.upper_threshold = calculate_thresholds(
+            self.bar,
+            self.bar.width()
+        )
 
 
 class BarFormat:
diff --git a/ascii_progress/spinner.py b/ascii_progress/spinner.py
index 17edecd..c3fc287 100644
--- a/ascii_progress/spinner.py
+++ b/ascii_progress/spinner.py
@@ -6,14 +6,55 @@
 # SPDX-License-Identifier: MIT
 
 import sys
-from typing import ContextManager, Iterator, TextIO, Sequence
+from typing import ContextManager, Iterator, TextIO, Sequence, TypeVar
 
 __all__ = (
-    "Spinner",
+    "SpinnerContext",
+    "Spinner"
 )
 
+T = TypeVar("T", bound="Spinner")
 
-class Spinner(ContextManager, Iterator[None]):
+
+class SpinnerContext(ContextManager[T]):
+    """context manager which handles exceptions while using the spinner"""
+
+    spinner: T
+
+    message: str
+
+    error: str
+
+    __slots__ = ("spinner", "message", "error")
+
+    def __init__(self, spinner: T, message: str, error: str) -> None:
+        self.spinner = spinner
+        self.message = message
+        self.error = error
+
+    def __eq__(self, other: object) -> bool:
+        if isinstance(other, SpinnerContext):
+            return self.spinner == other.spinner \
+                and self.message == other.message \
+                and self.error == other.error
+        return NotImplemented
+
+    def __enter__(self) -> T:
+        return self.spinner
+
+    def __exit__(self, type, value, traceback) -> bool:     # type: ignore
+        if type is None:
+            self.spinner.replace(self.message, end="\n")
+        elif type is KeyboardInterrupt:
+            # add 2 \b and 2 spaces to handle additional ^C
+            # add 2 additional spaces to make up for missing padding
+            self.spinner.replace("\b\b" + self.error, end="    \n")
+        else:
+            self.spinner.replace(self.error, end="\n")
+        return False    # we dont handle exceptions
+
+
+class Spinner(Iterator[None]):
     """class for creating a spinning animation"""
 
     frames: Sequence[str]
@@ -43,16 +84,6 @@ class Spinner(ContextManager, Iterator[None]):
                 and self.file is other.file
         return NotImplemented
 
-    def __enter__(self) -> "Spinner":
-        return self
-
-    def __exit__(self, type, value, traceback) -> bool:     # type: ignore
-        if type is KeyboardInterrupt:   # add 2 \b and 2 spaces to handle additional ^C
-            self.replace("\b\bKeyboardInterrupt", end="  \n")
-        else:
-            self.replace("Finished")
-        return False    # we dont handle exceptions
-
     def __iter__(self) -> "Spinner":
         return self
 
@@ -102,3 +133,7 @@ class Spinner(ContextManager, Iterator[None]):
         # pad message to fully overwrite old frame and add end
         self.file.write(message + " " * (len(self.current_frame) - len(message)) + end)
         self.file.flush()
+
+    def handle_exceptions(self: T, message: str, error: str) -> SpinnerContext[T]:
+        """return a context manager which replaces the spinner with message or error if a exceptions is raised"""
+        return SpinnerContext(self, message, error)


Test patch:
diff --git a/tests/test_bar.py b/tests/test_bar.py
index bfab94b..bc53278 100644
--- a/tests/test_bar.py
+++ b/tests/test_bar.py
@@ -5,8 +5,10 @@
 import unittest
 import unittest.mock
 from io import StringIO
+from typing import ContextManager
 from ascii_progress.bar import (
     LAZY_FORMATTER,
+    BarContext,
     Bar,
     ASCIIBar,
     ThresholdDecorator,
@@ -42,6 +44,58 @@ class TestLazyFormatter(unittest.TestCase):
         )
 
 
+class BarSpinnerContext(unittest.TestCase):
+    """tests for BarContext"""
+
+    def test_eq(self) -> None:
+        """test BarContext.__eq__"""
+        format = BarFormat(("-[", "]-"), ("  ", "=="), 5)
+        bar = ASCIIBar(format, file=StringIO())
+        contexts = [
+            BarContext(bar, "1", "2"),
+            BarContext(bar, "2", "2"),
+            BarContext(bar, "1", "1"),
+            BarContext(ASCIIBar(format, file=StringIO()), "1", "1")
+        ]
+        for context in contexts:
+            self.assertEqual(
+                [context],
+                [c for c in contexts if c == context]
+            )
+        self.assertNotEqual(contexts[0], 42)
+
+    def test_context(self) -> None:
+        """test SpinnerContext as context manager"""
+        output = StringIO()
+        bar = ASCIIBar(BarFormat(("[", "]"), (" ", "="), 3, format="{bar}"), file=output)
+        with BarContext(bar, "1", "2") as context:
+            self.assertIsInstance(context, Bar)
+            context.set_progress(34)
+            context.update()
+        with self.assertRaises(RuntimeError):
+            with BarContext(bar, "1", "2") as context:
+                context.set_progress(67)
+                context.update()
+                raise RuntimeError
+        with self.assertRaises(KeyboardInterrupt):
+            with BarContext(bar, "1", "2") as context:
+                context.set_progress(100)
+                context.update()
+                raise KeyboardInterrupt
+        self.assertEqual(
+            output.getvalue(),
+            "\b\b\b\b\b".join((
+                "[   ]",
+                "[=  ]",
+                "1    \n",
+                "[== ]",
+                "2    \n",
+                "[===]",
+                "\b\b2      \n"
+            ))
+        )
+
+
 class TestBar(unittest.TestCase):
     """Tests for Bar"""
 
@@ -76,18 +130,11 @@ class TestBar(unittest.TestCase):
                 break
         self.assertEqual(next(iterations), 11)
 
-    def test_enter(self) -> None:
-        """test Bar.__enter__ and __exit__"""
-        mock = unittest.mock.Mock(spec=Bar)
-        self.assertIs(Bar.__enter__(mock), mock)
-        self.assertFalse(Bar.__exit__(mock, None, None, None))
-        self.assertFalse(Bar.__exit__(mock, KeyboardInterrupt, KeyboardInterrupt(), None))
-        self.assertEqual(
-            mock.replace.call_args_list,
-            [
-                unittest.mock.call("Finished"),
-                unittest.mock.call("\b\bKeyboardInterrupt", end="  \n")
-            ]
+    def test_handle_exceptions(self) -> None:
+        """test Bar.handle_exceptions"""
+        self.assertIsInstance(
+            Bar.handle_exceptions(None, "", ""),
+            ContextManager
         )
 
 
diff --git a/tests/test_spinner.py b/tests/test_spinner.py
index 037d43c..0cffa5b 100644
--- a/tests/test_spinner.py
+++ b/tests/test_spinner.py
@@ -5,7 +5,48 @@
 import unittest
 import unittest.mock
 from io import StringIO
-from ascii_progress.spinner import Spinner
+from typing import ContextManager
+from ascii_progress.spinner import Spinner, SpinnerContext
+
+
+class TestSpinnerContext(unittest.TestCase):
+    """tests for SpinnerContext"""
+
+    def test_eq(self) -> None:
+        """test SpinnerContext.__eq__"""
+        spinner = Spinner("abc", StringIO())
+        contexts = [
+            SpinnerContext(spinner, "1", "2"),
+            SpinnerContext(spinner, "2", "2"),
+            SpinnerContext(spinner, "1", "1"),
+            SpinnerContext(Spinner("abc", StringIO()), "1", "2")
+        ]
+        for context in contexts:
+            self.assertEqual(
+                [context],
+                [c for c in contexts if c == context]
+            )
+        self.assertNotEqual(contexts[0], 42)
+
+    def test_context(self) -> None:
+        """test SpinnerContext as context manager"""
+        output = StringIO()
+        spinner = Spinner("abc", output)
+        with SpinnerContext(spinner, "1", "2") as context:
+            self.assertIsInstance(context, Spinner)
+            context.current_frame = "b"
+        with self.assertRaises(RuntimeError):
+            with SpinnerContext(spinner, "1", "2") as context:
+                context.current_frame = "c"
+                raise RuntimeError
+        with self.assertRaises(KeyboardInterrupt):
+            with SpinnerContext(spinner, "1", "2") as context:
+                context.current_frame = "a"
+                raise KeyboardInterrupt
+        self.assertEqual(
+            output.getvalue(),
+            "a\bb\b1\n\bc\b2\n\ba\b\b\b2    \n"
+        )
 
 
 class TestSpinner(unittest.TestCase):
@@ -31,27 +72,13 @@ class TestSpinner(unittest.TestCase):
             )
         self.assertNotEqual(spinners[0], 42)
 
-    def test_context(self) -> None:
-        """test Spinner __enter__ and __exit__"""
-        output = StringIO()
-        with Spinner("abc", output) as spinner:
-            self.assertIsInstance(spinner, Spinner)
-            self.assertEqual(output.getvalue(), "a")
-        self.assertEqual(output.getvalue(), "a\bFinished\n")
-        output.seek(0)
-        output.truncate(0)
-        with self.assertRaises(KeyboardInterrupt):
-            with Spinner("abc", output) as spinner:
-                raise KeyboardInterrupt
-        self.assertEqual(output.getvalue(), "a\b\b\bKeyboardInterrupt  \n")
-
     def test_iter(self) -> None:
         """test Spinner __iter__ and __exit__"""
         output = StringIO()
-        with Spinner("abc", output) as spinner:
-            for _ in zip(range(6), spinner):
-                pass
-        self.assertEqual(output.getvalue(), "a\bb\bc\ba\bb\bc\ba\bFinished\n")
+        spinner = Spinner("abc", output)
+        for _ in zip(range(6), spinner):
+            pass
+        self.assertEqual(output.getvalue(), "a\bb\bc\ba\bb\bc\ba")
 
     def test_with_padding(self) -> None:
         """test Spinner.with_padding"""
@@ -63,14 +90,14 @@ class TestSpinner(unittest.TestCase):
     def test_current_frame(self) -> None:
         """test Spinner.current_frame"""
         output = StringIO()
-        with Spinner("abc", output) as spinner:
-            self.assertEqual(spinner.current_frame, "a")
-            spinner.set_progress(1)
-            self.assertEqual(spinner.current_frame, "b")
-            spinner.current_frame = "a"
-            with self.assertRaises(ValueError):
-                spinner.current_frame = "x"
-        self.assertEqual(output.getvalue(), "a\bb\ba\bFinished\n")
+        spinner = Spinner("abc", output)
+        self.assertEqual(spinner.current_frame, "a")
+        spinner.set_progress(1)
+        self.assertEqual(spinner.current_frame, "b")
+        spinner.current_frame = "a"
+        with self.assertRaises(ValueError):
+            spinner.current_frame = "x"
+        self.assertEqual(output.getvalue(), "a\bb\ba")
 
     def test_update(self) -> None:
         """test Spinner.update"""
@@ -96,11 +123,13 @@ class TestSpinner(unittest.TestCase):
 
     def test_set_progress(self) -> None:
         """test Spinner.set_progress"""
-        output = StringIO()
-        with Spinner("abc", output) as spinner:
-            spinner.set_progress(2)
-            spinner.set_progress(4)
-        self.assertEqual(output.getvalue(), "a\bc\bb\bFinished\n")
+        output = unittest.mock.Mock(wraps=StringIO())
+        spinner = Spinner("abc", output)
+        output.reset_mock()
+        spinner.set_progress(2)
+        output.flush.assert_called()
+        spinner.set_progress(4)
+        self.assertEqual(output.getvalue(), "a\bc\bb")
 
     def test_replace(self) -> None:
         """test Spinner.replace"""
@@ -110,3 +139,10 @@ class TestSpinner(unittest.TestCase):
         Spinner(("testtest", ""), output).replace("test", "42")
         self.assertEqual(output.flush.call_count, 2)
         self.assertEqual(output.getvalue(), "testtest\b\b\b\b\b\b\b\btest    42")
+
+    def test_handle_exceptions(self) -> None:
+        """test Spinner.handle_exceptions"""
+        self.assertIsInstance(
+            Spinner("abc", file=StringIO()).handle_exceptions("", ""),
+            ContextManager
+        )