Repository: pallets/click
Pull request #2680 (2024-02-23 12:33:03)

Fix closing of callbacks on CLI exit
This PR forces the context to close itself on CLI exits.

This makes sure all callbacks registered by invoked options are properly called whenever a `ctx.exit()` action is triggered. Which has the effect of preventing state leaks.

These leaks cannot be witness in standard CLI operations, because on `ctx.exit()` call, the OS process that is running the Python CLI will be destroyed.

Still, in unittests, these state leaks will leads to flaky and non-deterministic tests, as the `pytest`, `tox` or any other test framework will keep using the same Python interpreter. To demonstrate this, I implemented a unit test using loggers as a fixture, in addition to a generic unit test.

Checklist:

- [x] Add tests that demonstrate the correct behavior of the change. Tests should fail without the change.
- [x] Add or update relevant docs, in the docs folder and in code.
- [x] Add an entry in `CHANGES.rst` summarizing the change and linking to the issue.
- [x] Add `.. versionchanged::` entries in any relevant code docs.
- [x] Run `pre-commit` hooks and fix any issues.
- [x] Run `pytest` and `tox`, no tests failed.


kdeldycke: Note that this might be related to: https://github.com/pallets/click/issues/1053

Patch:
diff --git a/CHANGES.rst b/CHANGES.rst
index cb4136a..59c5b58 100644
--- a/CHANGES.rst
+++ b/CHANGES.rst
@@ -45,6 +45,9 @@ Unreleased
     the ``mix_stderr`` parameter in ``CliRunner``. :issue:`2522` :pr:`2523`
 -   ``Option.show_envvar`` now also shows environment variable in error messages.
     :issue:`2695` :pr:`2696`
+-   ``Context.close`` will be called on exit. This results in all
+    ``Context.call_on_close`` callbacks and context managers added via
+    ``Context.with_resource`` to be closed on exit as well. :pr:`2680`
 
 
 Version 8.1.8
diff --git a/docs/advanced.rst b/docs/advanced.rst
index 54ced70..2911eac 100644
--- a/docs/advanced.rst
+++ b/docs/advanced.rst
@@ -494,3 +494,8 @@ cleanup function.
             db.record_use()
             db.save()
             db.close()
+
+
+.. versionchanged:: 8.2 ``Context.call_on_close`` and context managers registered
+    via ``Context.with_resource`` will be closed when the CLI exits. These were
+    previously not called on exit.
diff --git a/src/click/core.py b/src/click/core.py
index f17a8a8..968a10e 100644
--- a/src/click/core.py
+++ b/src/click/core.py
@@ -699,7 +699,13 @@ class Context:
         raise Abort()
 
     def exit(self, code: int = 0) -> t.NoReturn:
-        """Exits the application with a given exit code."""
+        """Exits the application with a given exit code.
+
+        .. versionchanged:: 8.2
+            Force closing of callbacks registered with
+            :meth:`call_on_close` before exiting the CLI.
+        """
+        self.close()
         raise Exit(code)
 
     def get_usage(self) -> str:


Test patch:
diff --git a/tests/test_context.py b/tests/test_context.py
index df8b497..5bd618b 100644
--- a/tests/test_context.py
+++ b/tests/test_context.py
@@ -1,9 +1,14 @@
+import logging
 from contextlib import contextmanager
 
 import pytest
 
 import click
+from click import Context
+from click import Option
+from click import Parameter
 from click.core import ParameterSource
+from click.decorators import help_option
 from click.decorators import pass_meta_key
 
 
@@ -237,6 +242,171 @@ def test_close_before_pop(runner):
     assert called == [True]
 
 
+def test_close_before_exit(runner):
+    called = []
+
+    @click.command()
+    @click.pass_context
+    def cli(ctx):
+        ctx.obj = "test"
+
+        @ctx.call_on_close
+        def foo():
+            assert click.get_current_context().obj == "test"
+            called.append(True)
+
+        ctx.exit()
+
+        click.echo("aha!")
+
+    result = runner.invoke(cli, [])
+    assert not result.exception
+    assert not result.output
+    assert called == [True]
+
+
+@pytest.mark.parametrize(
+    ("cli_args", "expect"),
+    [
+        pytest.param(
+            ("--option-with-callback", "--force-exit"),
+            ["ExitingOption", "NonExitingOption"],
+            id="natural_order",
+        ),
+        pytest.param(
+            ("--force-exit", "--option-with-callback"),
+            ["ExitingOption"],
+            id="eagerness_precedence",
+        ),
+    ],
+)
+def test_multiple_eager_callbacks(runner, cli_args, expect):
+    """Checks all callbacks are called on exit, even the nasty ones hidden within
+    callbacks.
+
+    Also checks the order in which they're called.
+    """
+    # Keeps track of callback calls.
+    called = []
+
+    class NonExitingOption(Option):
+        def reset_state(self):
+            called.append(self.__class__.__name__)
+
+        def set_state(self, ctx: Context, param: Parameter, value: str) -> str:
+            ctx.call_on_close(self.reset_state)
+            return value
+
+        def __init__(self, *args, **kwargs) -> None:
+            kwargs.setdefault("expose_value", False)
+            kwargs.setdefault("callback", self.set_state)
+            super().__init__(*args, **kwargs)
+
+    class ExitingOption(NonExitingOption):
+        def set_state(self, ctx: Context, param: Parameter, value: str) -> str:
+            value = super().set_state(ctx, param, value)
+            ctx.exit()
+            return value
+
+    @click.command()
+    @click.option("--option-with-callback", is_eager=True, cls=NonExitingOption)
+    @click.option("--force-exit", is_eager=True, cls=ExitingOption)
+    def cli():
+        click.echo("This will never be printed as we forced exit via --force-exit")
+
+    result = runner.invoke(cli, cli_args)
+    assert not result.exception
+    assert not result.output
+
+    assert called == expect
+
+
+def test_no_state_leaks(runner):
+    """Demonstrate state leaks with a specific case of the generic test above.
+
+    Use a logger as a real-world example of a common fixture which, due to its global
+    nature, can leak state if not clean-up properly in a callback.
+    """
+    # Keeps track of callback calls.
+    called = []
+
+    class DebugLoggerOption(Option):
+        """A custom option to set the name of the debug logger."""
+
+        logger_name: str
+        """The ID of the logger to use."""
+
+        def reset_loggers(self):
+            """Forces logger managed by the option to be reset to the default level."""
+            logger = logging.getLogger(self.logger_name)
+            logger.setLevel(logging.NOTSET)
+
+            # Logger has been properly reset to its initial state.
+            assert logger.level == logging.NOTSET
+            assert logger.getEffectiveLevel() == logging.WARNING
+
+            called.append(True)
+
+        def set_level(self, ctx: Context, param: Parameter, value: str) -> None:
+            """Set the logger to DEBUG level."""
+            # Keep the logger name around so we can reset it later when winding down
+            # the option.
+            self.logger_name = value
+
+            # Get the global logger object.
+            logger = logging.getLogger(self.logger_name)
+
+            # Check pre-conditions: new logger is not set, but inherits its level from
+            # default <root> logger. That's the exact same state we are expecting our
+            # logger to be in after being messed with by the CLI.
+            assert logger.level == logging.NOTSET
+            assert logger.getEffectiveLevel() == logging.WARNING
+
+            logger.setLevel(logging.DEBUG)
+            ctx.call_on_close(self.reset_loggers)
+            return value
+
+        def __init__(self, *args, **kwargs) -> None:
+            kwargs.setdefault("callback", self.set_level)
+            super().__init__(*args, **kwargs)
+
+    @click.command()
+    @click.option("--debug-logger-name", is_eager=True, cls=DebugLoggerOption)
+    @help_option()
+    @click.pass_context
+    def messing_with_logger(ctx, debug_logger_name):
+        # Introspect context to make sure logger name are aligned.
+        assert debug_logger_name == ctx.command.params[0].logger_name
+
+        logger = logging.getLogger(debug_logger_name)
+
+        # Logger's level has been properly set to DEBUG by DebugLoggerOption.
+        assert logger.level == logging.DEBUG
+        assert logger.getEffectiveLevel() == logging.DEBUG
+
+        logger.debug("Blah blah blah")
+
+        ctx.exit()
+
+        click.echo("This will never be printed as we exited early")
+
+    # Call the CLI to mess with the custom logger.
+    result = runner.invoke(
+        messing_with_logger, ["--debug-logger-name", "my_logger", "--help"]
+    )
+
+    assert called == [True]
+
+    # Check the custom logger has been reverted to it initial state by the option
+    # callback after being messed with by the CLI.
+    logger = logging.getLogger("my_logger")
+    assert logger.level == logging.NOTSET
+    assert logger.getEffectiveLevel() == logging.WARNING
+
+    assert not result.exception
+    assert result.output.startswith("Usage: messing-with-logger [OPTIONS]")
+
+
 def test_with_resource():
     @contextmanager
     def manager():