Repository: pallets/click
Pull request #2421 (2022-12-24 18:12:05)

Fix return value and type signature of `shell_completion.add_completion_class` function
Decorators should return a possibly modfiied version of the original type. Hence, add_completion_class should return cls rather than merely None. This conforms to Python conventions, and allows type checking to work properly.

Checklist:

- [x] Add tests that demonstrate the correct behavior of the change. Tests should fail without the change.
- [ ] 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.
- [ ] 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.

alexreg: Pre-commit check failures seem to concern pre-existing issues.
alexreg: Hi @davidism. Is this okay for you?
saroad2: Hey @alexreg , thanks for the PR.

While this PR is pretty straight-forward, you didn't include any tests to verify it.
Can you add some tests to check that everything works smoothly?
alexreg: @saroad2 Sure. I added a couple of tests for the return value.
saroad2: Hey @alexreg , thank you for reposing my review. I'm glad I was able to introduce you to new concepts.

I've added some small tweaks to your PR, mainly a small refactoring for the fixture, elaboration on the steps of the tests you wrote, and add one more test about using `add_completion_class` with the `name` parameter.

Happily merging the PR and hope to see your contributions again in the future!

Patch:
diff --git a/CHANGES.rst b/CHANGES.rst
index c60c12a..7b1e5b3 100644
--- a/CHANGES.rst
+++ b/CHANGES.rst
@@ -7,6 +7,7 @@ Unreleased
 
 -   Improve type hinting for decorators and give all generic types parameters.
     :issue:`2398`
+-   Fix return value and type signature of `shell_completion.add_completion_class` function. :pr:`2421`
 
 
 Version 8.1.3
diff --git a/src/click/shell_completion.py b/src/click/shell_completion.py
index c17a8e6..3b3a1df 100644
--- a/src/click/shell_completion.py
+++ b/src/click/shell_completion.py
@@ -389,6 +389,9 @@ class FishComplete(ShellComplete):
         return f"{item.type},{item.value}"
 
 
+ShellCompleteType = t.TypeVar("ShellCompleteType", bound=t.Type[ShellComplete])
+
+
 _available_shells: t.Dict[str, t.Type[ShellComplete]] = {
     "bash": BashComplete,
     "fish": FishComplete,
@@ -397,8 +400,8 @@ _available_shells: t.Dict[str, t.Type[ShellComplete]] = {
 
 
 def add_completion_class(
-    cls: t.Type[ShellComplete], name: t.Optional[str] = None
-) -> None:
+    cls: ShellCompleteType, name: t.Optional[str] = None
+) -> ShellCompleteType:
     """Register a :class:`ShellComplete` subclass under the given name.
     The name will be provided by the completion instruction environment
     variable during completion.
@@ -413,6 +416,8 @@ def add_completion_class(
 
     _available_shells[name] = cls
 
+    return cls
+
 
 def get_completion_class(shell: str) -> t.Optional[t.Type[ShellComplete]]:
     """Look up a registered :class:`ShellComplete` subclass by the name


Test patch:
diff --git a/tests/test_shell_completion.py b/tests/test_shell_completion.py
index 57729f5..3e7a2bd 100644
--- a/tests/test_shell_completion.py
+++ b/tests/test_shell_completion.py
@@ -1,9 +1,11 @@
 import pytest
 
+import click.shell_completion
 from click.core import Argument
 from click.core import Command
 from click.core import Group
 from click.core import Option
+from click.shell_completion import add_completion_class
 from click.shell_completion import CompletionItem
 from click.shell_completion import ShellComplete
 from click.types import Choice
@@ -342,3 +344,73 @@ def test_choice_case_sensitive(value, expect):
     )
     completions = _get_words(cli, ["-a"], "a")
     assert completions == expect
+
+
+@pytest.fixture()
+def _restore_available_shells(tmpdir):
+    prev_available_shells = click.shell_completion._available_shells.copy()
+    click.shell_completion._available_shells.clear()
+    yield
+    click.shell_completion._available_shells.clear()
+    click.shell_completion._available_shells.update(prev_available_shells)
+
+
+@pytest.mark.usefixtures("_restore_available_shells")
+def test_add_completion_class():
+    # At first, "mysh" is not in available shells
+    assert "mysh" not in click.shell_completion._available_shells
+
+    class MyshComplete(ShellComplete):
+        name = "mysh"
+        source_template = "dummy source"
+
+    # "mysh" still not in available shells because it is not registered
+    assert "mysh" not in click.shell_completion._available_shells
+
+    # Adding a completion class should return that class
+    assert add_completion_class(MyshComplete) is MyshComplete
+
+    # Now, "mysh" is finally in available shells
+    assert "mysh" in click.shell_completion._available_shells
+    assert click.shell_completion._available_shells["mysh"] is MyshComplete
+
+
+@pytest.mark.usefixtures("_restore_available_shells")
+def test_add_completion_class_with_name():
+    # At first, "mysh" is not in available shells
+    assert "mysh" not in click.shell_completion._available_shells
+    assert "not_mysh" not in click.shell_completion._available_shells
+
+    class MyshComplete(ShellComplete):
+        name = "not_mysh"
+        source_template = "dummy source"
+
+    # "mysh" and "not_mysh" are still not in available shells because
+    # it is not registered yet
+    assert "mysh" not in click.shell_completion._available_shells
+    assert "not_mysh" not in click.shell_completion._available_shells
+
+    # Adding a completion class should return that class.
+    # Because we are using the "name" parameter, the name isn't taken
+    # from the class.
+    assert add_completion_class(MyshComplete, name="mysh") is MyshComplete
+
+    # Now, "mysh" is finally in available shells
+    assert "mysh" in click.shell_completion._available_shells
+    assert "not_mysh" not in click.shell_completion._available_shells
+    assert click.shell_completion._available_shells["mysh"] is MyshComplete
+
+
+@pytest.mark.usefixtures("_restore_available_shells")
+def test_add_completion_class_decorator():
+    # At first, "mysh" is not in available shells
+    assert "mysh" not in click.shell_completion._available_shells
+
+    @add_completion_class
+    class MyshComplete(ShellComplete):
+        name = "mysh"
+        source_template = "dummy source"
+
+    # Using `add_completion_class` as a decorator adds the new shell immediately
+    assert "mysh" in click.shell_completion._available_shells
+    assert click.shell_completion._available_shells["mysh"] is MyshComplete