address interlinked issues #36, #35 and #27

bckohan · bckohan · commit f435e5e9e3bd · 2024-03-04T23:13:21.000-08:00
diff --git a/django_typer/__init__.py b/django_typer/__init__.py
@@ -107,7 +107,7 @@
 )
 from .utils import _command_context, traceback_config, with_typehint
 
-VERSION = (1, 0, 1)
+VERSION = (1, 0, 2)
 
 __title__ = "Django Typer"
 __version__ = ".".join(str(i) for i in VERSION)
@@ -298,6 +298,7 @@ class _ParsedArgs(SimpleNamespace):  # pylint: disable=too-few-public-methods
     def __init__(self, args: t.Sequence[t.Any], **kwargs):
         super().__init__(**kwargs)
         self.args = args
+        self.traceback = kwargs.get("traceback", TyperCommand._traceback)
 
     def _get_kwargs(self):
         return {"args": self.args, **COMMON_DEFAULTS}
@@ -1451,23 +1452,11 @@ def option_strings(self) -> t.List[str]:
     prog_name: str
     subcommand: str
 
-    missing_args_message: str = "Missing parameter: {parameter}"
-    called_from_command_line: t.Optional[bool] = None
-
-    def __init__(
-        self,
-        django_command: "TyperCommand",
-        prog_name,
-        subcommand,
-        missing_args_message: str = missing_args_message,
-        called_from_command_line: t.Optional[bool] = None,
-    ):
+    def __init__(self, django_command: "TyperCommand", prog_name, subcommand):
         self._actions = []
         self.django_command = django_command
         self.prog_name = prog_name
         self.subcommand = subcommand
-        self.missing_args_message = missing_args_message
-        self.called_from_command_line = called_from_command_line
 
         def populate_params(node: CommandNode) -> None:
             for param in node.click_command.params:
@@ -1505,29 +1494,18 @@ def parse_args(self, args=None, namespace=None) -> _ParsedArgs:
             (this is ignored for TyperCommands but is required by the django
             base class)
         """
-        try:
+        with self.django_command:
             cmd = get_typer_command(self.django_command.typer_app)
             with cmd.make_context(
                 info_name=f"{self.prog_name} {self.subcommand}",
                 django_command=self.django_command,
                 args=list(args or []),
             ) as ctx:
-                return _ParsedArgs(args=args or [], **{**COMMON_DEFAULTS, **ctx.params})
-        except click.exceptions.Exit:
-            sys.exit()
-        except (click.exceptions.UsageError, CommandError) as arg_err:
-            if self.called_from_command_line:
-                err_msg = (
-                    _(self.missing_args_message).format(
-                        parameter=getattr(getattr(arg_err, "param", None), "name", "")
-                    )
-                    if isinstance(arg_err, click.exceptions.MissingParameter)
-                    else str(arg_err)
+                common = {**COMMON_DEFAULTS, **ctx.params}
+                self.django_command._traceback = common.get(
+                    "traceback", self.django_command._traceback
                 )
-                self.print_help()
-                self.django_command.stderr.write(err_msg)
-                sys.exit(1)
-            raise CommandError(str(arg_err)) from arg_err
+                return _ParsedArgs(args=args or [], **common)
 
     def add_argument(self, *args, **kwargs):
         """
@@ -1686,27 +1664,57 @@ def command2(self, option: t.Optional[str] = None):
     # they can use the type from django_typer.types.Verbosity
     suppressed_base_arguments = {"verbosity"}
 
+    missing_args_message = "Missing parameter: {parameter}"
+
     typer_app: Typer
     no_color: bool = False
     force_color: bool = False
     _num_commands: int = 0
     _has_callback: bool = False
     _root_groups: int = 0
     _handle: t.Callable[..., t.Any]
+    _traceback: bool = False
 
     command_tree: CommandNode
 
     @property
-    def name(self):
+    def name(self) -> str:
         """The name of the django command"""
-        return self.typer_app.info.name
+        return self.typer_app.info.name or self.__module__.rsplit(".", maxsplit=1)[-1]
 
     def __enter__(self):
         _command_context.__dict__.setdefault("stack", []).append(self)
         return self
 
     def __exit__(self, exc_type, exc_val, exc_tb):
         _command_context.stack.pop()
+        if isinstance(exc_val, click.exceptions.Exit):
+            sys.exit(exc_val.exit_code)
+        if isinstance(exc_val, click.exceptions.UsageError):
+            err_msg = (
+                _(self.missing_args_message).format(
+                    parameter=getattr(getattr(exc_val, "param", None), "name", "")
+                )
+                if isinstance(exc_val, click.exceptions.MissingParameter)
+                else str(exc_val)
+            )
+
+            # we might be in a subcommand - so make sure we pull that help out
+            # by walking up the context tree until we're at root
+            cmd_pth: t.List[str] = []
+            ctx = exc_val.ctx
+            while ctx and ctx.parent:
+                if ctx.info_name:
+                    cmd_pth.insert(0, ctx.info_name)
+                ctx = ctx.parent
+            if (
+                getattr(self, "_called_from_command_line", False)
+                and not self._traceback
+            ):
+                self.print_help(sys.argv[0], self.name, *cmd_pth)
+                self.stderr.write(err_msg)
+                sys.exit(1)
+            raise CommandError(str(exc_val)) from exc_val
 
     def __init__(
         self,
@@ -1816,17 +1824,7 @@ def create_parser(  # pyright: ignore[reportIncompatibleMethodOverride]
         :param subcommand: the name of the django command
         """
         with self:
-            return TyperParser(
-                self,
-                prog_name,
-                subcommand,
-                missing_args_message=getattr(
-                    self, "missing_args_message", TyperParser.missing_args_message
-                ),
-                called_from_command_line=getattr(
-                    self, "_called_from_command_line", None
-                ),
-            )
+            return TyperParser(self, prog_name, subcommand)
 
     def print_help(self, prog_name: str, subcommand: str, *cmd_path: str):
         """
@@ -1936,8 +1934,9 @@ def execute(self, *args, **options):
             self.no_color = options["no_color"]
         if options.get("force_color", None) is not None:
             self.force_color = options["force_color"]
-        with self:
-            result = super().execute(*args, **options)
-        self.no_color = no_color
-        self.force_color = force_color
-        return result
+        try:
+            with self:
+                return super().execute(*args, **options)
+        finally:
+            self.no_color = no_color
+            self.force_color = force_color
diff --git a/django_typer/apps.py b/django_typer/apps.py
@@ -23,8 +23,11 @@
 rich: t.Union[ModuleType, None] = None
 
 try:
+    import sys
+
     import rich
     from rich import traceback
+    from typer import main as typer_main
 
     tb_config = traceback_config()
     if rich and isinstance(tb_config, dict) and not tb_config.get("no_install", False):
@@ -52,6 +55,10 @@
                 if param in set(inspect.signature(traceback.install).parameters.keys())
             },
         )
+        # typer installs its own exception hook and it falls back to the sys hook - depending
+        # on when typer was imported it may have the original fallback system hook or our
+        # installed rich one - we patch it here to make sure!
+        typer_main._original_except_hook = sys.excepthook
 except ImportError:
     pass
 
diff --git a/django_typer/tests/test_app/management/commands/tb.py b/django_typer/tests/test_app/management/commands/tb.py
@@ -0,0 +1,26 @@
+import sys
+
+import click
+from django.core.management.base import CommandError
+
+from django_typer import TyperCommand, command
+
+
+class Command(TyperCommand):
+
+    help = "Test traceback scenarios."
+
+    @command()
+    def error(self, throw_command: bool = False, throw_other: bool = False):
+        if throw_command and throw_other:
+            raise click.exceptions.UsageError(
+                f"--throw-command and --throw-other are mutually exclusive."
+            )
+        if throw_command:
+            raise CommandError(f"This command failed!")
+        if throw_other:
+            raise RuntimeError(f"This command threw an unexpected error!")
+
+    @command()
+    def exit(self, code: int = 0):
+        sys.exit(code)
diff --git a/django_typer/tests/test_app/management/commands/vanilla.py b/django_typer/tests/test_app/management/commands/vanilla.py
@@ -1,12 +1,11 @@
-import json
-from pprint import pprint
+import sys
 
 from django.core.management.base import BaseCommand, CommandParser
 
 
 class Command(BaseCommand):
     def add_arguments(self, parser: CommandParser) -> None:
-        parser.add_argument("--name", type=str)
+        parser.add_argument("--exit-code", dest="exit_code", type=int, default=0)
 
     def handle(self, **options):
-        pprint(json.dumps(options))
+        raise sys.exit(options.get("exit_code", 0))
diff --git a/django_typer/tests/tests.py b/django_typer/tests/tests.py
@@ -77,17 +77,17 @@ def run_command(command, *args, parse_json=True) -> Tuple[str, str]:
 
         # Check the return code to ensure the script ran successfully
         if result.returncode != 0:
-            return result.stdout, result.stderr
+            return result.stdout, result.stderr, result.returncode
 
         # Parse the output
         if result.stdout:
             if parse_json:
                 try:
-                    return json.loads(result.stdout), result.stderr
+                    return json.loads(result.stdout), result.stderr, result.returncode
                 except json.JSONDecodeError:
-                    return result.stdout, result.stderr
-            return result.stdout, result.stderr
-        return result.stdout, result.stderr
+                    return result.stdout, result.stderr, result.returncode
+            return result.stdout, result.stderr, result.returncode
+        return result.stdout, result.stderr, result.returncode
     finally:
         os.chdir(cwd)
 
@@ -1271,7 +1271,8 @@ def test_command_line(self, settings=None):
                 ("hey! " * 5).strip(),
             )
         else:
-            self.assertIn("UsageError", result[1])
+            self.assertIn("Usage: ./manage.py groups echo [OPTIONS] MESSAGE", result[0])
+            self.assertIn("Got unexpected extra argument (5)", result[1])
             with self.assertRaises(TypeError):
                 call_command("groups", "echo", "hey!", echoes=5)
             with self.assertRaises(TypeError):
@@ -1415,13 +1416,16 @@ def test_command_line(self, settings=None):
             "groups", *settings, "string", "annamontes", "case", "upper", "4", "9"
         )
         if override:
-            result = result[1].strip()
-            self.assertIn("UsageError", result)
+            self.assertIn(
+                "Usage: ./manage.py groups string STRING case upper [OPTIONS]",
+                result[0],
+            )
+            self.assertIn("Got unexpected extra arguments (4 9)", result[1].strip())
             grp_cmd.string("annamontes")
             with self.assertRaises(TypeError):
                 self.assertEqual(grp_cmd.upper(4, 9), "annaMONTEs")
 
-            with self.assertRaises(UsageError):
+            with self.assertRaises(CommandError):
                 self.assertEqual(
                     call_command(
                         "groups", "string", "annamontes", "case", "upper", "4", "9"
@@ -1451,8 +1455,12 @@ def test_command_line(self, settings=None):
             grp_cmd.string(" emmatc  ")
             self.assertEqual(grp_cmd.strip(), "emmatc")
         else:
-            self.assertIn("UsageError", result[1])
-            with self.assertRaises(UsageError):
+            self.assertIn(
+                "Usage: ./manage.py groups string [OPTIONS] STRING COMMAND [ARGS]",
+                result[0],
+            )
+            self.assertIn("No such command 'strip'.", result[1])
+            with self.assertRaises(CommandError):
                 self.assertEqual(
                     call_command("groups", "string", " emmatc  ", "strip"), "emmatc"
                 )
@@ -2615,3 +2623,85 @@ def test_custom_fallback(self):
             "shell ",
         )[0]
         self.assertTrue("shell " in result)
+
+
+class TracebackTests(TestCase):
+    """
+    Tests that show CommandErrors and UsageErrors do not result in tracebacks unless --traceback is set.
+
+    Also make sure that sys.exit is not called when not run from the terminal
+    (i.e. in get_command invocation or call_command).
+    """
+
+    def test_usage_error_no_tb(self):
+        stdout, stderr, retcode = run_command("tb", "wrong")
+        self.assertTrue("Usage: ./manage.py tb [OPTIONS] COMMAND [ARGS]" in stdout)
+        self.assertTrue("No such command" in stderr)
+        self.assertTrue(retcode > 0)
+
+        stdout, stderr, retcode = run_command("tb", "error", "wrong")
+        self.assertTrue("Usage: ./manage.py tb error [OPTIONS]" in stdout)
+        self.assertTrue("Got unexpected extra argument" in stderr)
+        self.assertTrue(retcode > 0)
+
+        with self.assertRaises(CommandError):
+            call_command("tb", "wrong")
+
+        with self.assertRaises(CommandError):
+            call_command("tb", "error", "wrong")
+
+    def test_usage_error_with_tb_if_requested(self):
+
+        stdout, stderr, retcode = run_command("tb", "--traceback", "wrong")
+        self.assertFalse(stdout.strip())
+        self.assertTrue("Traceback" in stderr)
+        if rich_installed:
+            self.assertTrue("───── locals ─────" in stderr)
+        else:
+            self.assertFalse("───── locals ─────" in stderr)
+        self.assertTrue("No such command 'wrong'" in stderr)
+        self.assertTrue(retcode > 0)
+
+        stdout, stderr, retcode = run_command("tb", "--traceback", "error", "wrong")
+        self.assertFalse(stdout.strip())
+        self.assertTrue("Traceback" in stderr)
+        if rich_installed:
+            self.assertTrue("───── locals ─────" in stderr)
+        else:
+            self.assertFalse("───── locals ─────" in stderr)
+        self.assertFalse(stdout.strip())
+        self.assertTrue("Got unexpected extra argument" in stderr)
+        self.assertTrue(retcode > 0)
+
+        with self.assertRaises(CommandError):
+            call_command("tb", "--traceback", "wrong")
+
+        with self.assertRaises(CommandError):
+            call_command("tb", "--traceback", "error", "wrong")
+
+    def test_click_exception_retcodes_honored(self):
+
+        self.assertEqual(run_command("vanilla")[2], 0)
+        self.assertEqual(run_command("vanilla", "--exit-code=2")[2], 2)
+
+        self.assertEqual(run_command("tb", "exit")[2], 0)
+        self.assertEqual(run_command("tb", "exit", "--code=2")[2], 2)
+
+    def test_exit_on_call(self):
+        with self.assertRaises(SystemExit):
+            call_command("vanilla", "--help")
+
+        with self.assertRaises(SystemExit):
+            call_command("vanilla", "--exit-code", "0")
+
+        with self.assertRaises(SystemExit):
+            call_command("vanilla", "--exit-code", "1")
+
+        with self.assertRaises(SystemExit):
+            call_command("tb", "--help")
+
+        with self.assertRaises(SystemExit):
+            call_command("tb", "exit")
+
+        with self.assertRaises(SystemExit):
+            call_command("tb", "exit", "--code=1")
diff --git a/doc/source/changelog.rst b/doc/source/changelog.rst
diff --git a/pyproject.toml b/pyproject.toml
