Allow the exception handling in TracebackHandler to be customised.

Author: domdfcoding

consolekit/tracebacks.py

@@ -5,6 +5,7 @@
 Functions for handling exceptions and their tracebacks.
 
 .. versionadded:: 1.0.0
+.. latex:vspace:: -10px
 """
 #
 #  Copyright © 2020-2021 Dominic Davis-Foster <[email protected]>
@@ -30,12 +31,17 @@
 
 # stdlib
 import contextlib
-from typing import Callable, ContextManager, Type, TypeVar
+import sys
+from typing import TYPE_CHECKING, Callable, ContextManager, List, Type, TypeVar, Union
 
 # 3rd party
 import click
 from domdf_python_tools.compat import nullcontext
 
+if sys.version_info >= (3, 7) or TYPE_CHECKING:
+	# stdlib
+	from typing import NoReturn
+
 __all__ = ["TracebackHandler", "handle_tracebacks", "traceback_handler", "traceback_option"]
 
 _C = TypeVar("_C", bound=click.Command)
@@ -65,38 +71,65 @@ class TracebackHandler:
 	is the name of the exception class to handle.
 
 	.. versionadded:: 1.0.0
+
+	:param exception: The exception to raise after handling the traceback.
+		If not running within a click command or group you'll likely
+		want to set this to :exc:`SystemExit(1) <SystemExit>`.
+	:default exception: :class:`click.Abort() <click.Abort>`
+
+	.. versionchanged:: 1.4.0  Added the ``exception`` argument.
+
 	.. seealso:: :func:`~.handle_tracebacks`.
 	.. latex:clearpage::
 	"""  # noqa: D400
 
-	def handle_EOFError(self, e: EOFError) -> bool:  # noqa: D102
-		raise e
+	exception: Exception
+	"""
+	The exception to raise after handling the traceback.
 
-	def handle_KeyboardInterrupt(self, e: KeyboardInterrupt) -> bool:  # noqa: D102
-		raise e
+	.. versionadded:: 1.4.0
+	"""
 
-	def handle_ClickException(self, e: click.ClickException) -> bool:  # noqa: D102
-		raise e
+	def __init__(self, exception: BaseException = click.Abort()):
+		self.exception: BaseException = exception  # type: ignore[assignment]
 
-	def handle_Abort(self, e: click.Abort) -> bool:  # noqa: D102
-		raise e
+	def abort(self, msg: Union[str, List[str]]) -> "NoReturn":
+		"""
+		Abort the current process by calling ``self.exception``.
+
+		.. versionadded:: 1.4.0
+
+		:param msg: The message to write to stderr before raising the exception.
+		"""
 
-	def handle_SystemExit(self, e: SystemExit) -> bool:  # noqa: D102
+		if isinstance(msg, str):
+			msg = ''.join(msg)
+
+		if msg:
+			click.echo(msg, err=True, color=False)
+
+		raise self.exception
+
+	def handle_EOFError(self, e: EOFError) -> "NoReturn":  # noqa: D102
 		raise e
 
-	def handle_FileNotFoundError(self, e: FileNotFoundError) -> bool:  # noqa: D102
+	def handle_KeyboardInterrupt(self, e: KeyboardInterrupt) -> "NoReturn":  # noqa: D102
+		raise e
 
-		# this package
-		from consolekit.utils import abort
+	def handle_ClickException(self, e: click.ClickException) -> "NoReturn":  # noqa: D102
+		raise e
 
-		raise abort(f"File Not Found: {e}", colour=False)
+	def handle_Abort(self, e: click.Abort) -> "NoReturn":  # noqa: D102
+		raise e
 
-	def handle_FileExistsError(self, e: FileExistsError) -> bool:  # noqa: D102
+	def handle_SystemExit(self, e: SystemExit) -> "NoReturn":  # noqa: D102
+		raise e
 
-		# this package
-		from consolekit.utils import abort
+	def handle_FileNotFoundError(self, e: FileNotFoundError) -> "NoReturn":  # noqa: D102
+		self.abort(f"File Not Found: {e}")
 
-		raise abort(f"File Exists: {e}", colour=False)
+	def handle_FileExistsError(self, e: FileExistsError) -> "NoReturn":  # noqa: D102
+		self.abort(f"File Exists: {e}")
 
 	def handle(self, e: BaseException) -> bool:
 		"""
@@ -117,7 +150,7 @@ def handle(self, e: BaseException) -> bool:
 			if hasattr(self, f"handle_{base.__name__}"):
 				return getattr(self, f"handle_{base.__name__}")(e)
 
-		raise abort(f"An error occurred: {e}", colour=False)
+		self.abort(f"An error occurred: {e}")
 
 	@contextlib.contextmanager
 	def __call__(self):
@@ -171,8 +204,9 @@ def handle_tracebacks(
 		In either case the program execution will stop on error.
 	:param cls: The class to use to handle the tracebacks.
 
-	.. versionchanged:: 1.0.0  Added the ``cls`` parameter.
+	:rtype:
 
+	.. versionchanged:: 1.0.0  Added the ``cls`` parameter.
 	.. seealso:: :func:`~.traceback_handler` and :class:`~.TracebackHandler`
 	"""
 

tests/test_tracebacks.py

@@ -26,11 +26,13 @@
 				]
 		)
 contextmanagers = pytest.mark.parametrize(
-		"contextmanager",
+		"contextmanager, exit_code",
 		[
-				pytest.param(handle_tracebacks, id="handle_tracebacks"),
-				pytest.param(traceback_handler, id="traceback_handler"),
-				pytest.param(TracebackHandler(), id="TracebackHandler"),
+				pytest.param(handle_tracebacks, 1, id="handle_tracebacks"),
+				pytest.param(traceback_handler, 1, id="traceback_handler"),
+				pytest.param(TracebackHandler(), 1, id="TracebackHandler"),
+				pytest.param(TracebackHandler(SystemExit(1)), 1, id="TracebackHandler_SystemExit"),
+				pytest.param(TracebackHandler(SystemExit(2)), 2, id="TracebackHandler_SystemExit_2"),
 				]
 		)
 
@@ -40,6 +42,7 @@
 def test_handle_tracebacks(
 		exception,
 		contextmanager: Callable[..., ContextManager],
+		exit_code: int,
 		file_regression,
 		cli_runner: CliRunner,
 		):
@@ -52,7 +55,7 @@ def demo():
 
 	result: Result = cli_runner.invoke(demo, catch_exceptions=False)
 	result.check_stdout(file_regression)
-	assert result.exit_code == 1
+	assert result.exit_code == exit_code
 
 
 @exceptions
@@ -74,6 +77,7 @@ def test_handle_tracebacks_ignored_exceptions_click(
 		exception: Exception,
 		contextmanager: Callable[..., ContextManager],
 		cli_runner: CliRunner,
+		exit_code: int,
 		):
 
 	@click.command()
@@ -93,6 +97,7 @@ def demo():
 def test_handle_tracebacks_ignored_exceptions(
 		exception: Type[Exception],
 		contextmanager: Callable[..., ContextManager],
+		exit_code: int,
 		):
 
 	with pytest.raises(exception):  # noqa: PT012
@@ -140,7 +145,8 @@ def test_handle_tracebacks_ignored_click(
 		advanced_file_regression: AdvancedFileRegressionFixture,
 		code: int,
 		cli_runner: CliRunner,
-		click_version: str
+		click_version: str,
+		exit_code: int,
 		):
 
 	@click.command()
@@ -151,6 +157,10 @@ def demo():
 
 	result = cli_runner.invoke(demo, catch_exceptions=False)
 	result.check_stdout(advanced_file_regression)
+
+	# if code == 1 and exit_code == 2:
+	# 	code = 2
+
 	assert result.exit_code == code
 
 

tests/test_tracebacks_/test_handle_tracebacks_TracebackHandler_SystemExit_2_Exception_.txt

@@ -0,0 +1 @@
+An error occurred: Something's awry!

tests/test_tracebacks_/test_handle_tracebacks_TracebackHandler_SystemExit_2_FileExistsError_.txt

@@ -0,0 +1 @@
+File Exists: foo.txt

tests/test_tracebacks_/test_handle_tracebacks_TracebackHandler_SystemExit_2_FileNotFoundError_.txt

@@ -0,0 +1 @@
+File Not Found: foo.txt

tests/test_tracebacks_/test_handle_tracebacks_TracebackHandler_SystemExit_2_NameError_.txt

@@ -0,0 +1 @@
+An error occurred: name 'hello' is not defined

tests/test_tracebacks_/test_handle_tracebacks_TracebackHandler_SystemExit_2_SyntaxError_.txt

@@ -0,0 +1 @@
+An error occurred: invalid syntax

tests/test_tracebacks_/test_handle_tracebacks_TracebackHandler_SystemExit_2_TypeError_.txt

@@ -0,0 +1 @@
+An error occurred: Expected type int, got type str

tests/test_tracebacks_/test_handle_tracebacks_TracebackHandler_SystemExit_2_ValueError_.txt

@@ -0,0 +1 @@
+An error occurred: 'age' must be >= 0

tests/test_tracebacks_/test_handle_tracebacks_TracebackHandler_SystemExit_Exception_.txt

@@ -0,0 +1 @@
+An error occurred: Something's awry!