querysets can be passed to model object completers, fixes #96

Author: bckohan

django_typer/__init__.py

@@ -47,7 +47,7 @@
     model_parser_completer,  # noqa: F401
 )
 
-VERSION = (2, 1, 3)
+VERSION = (2, 2, 0)
 
 __title__ = "Django Typer"
 __version__ = ".".join(str(i) for i in VERSION)

django_typer/completers.py

@@ -16,6 +16,7 @@
 
 # pylint: disable=line-too-long
 
+import inspect
 import os
 import pkgutil
 import sys
@@ -38,12 +39,15 @@
     FloatField,
     GenericIPAddressField,
     IntegerField,
+    Manager,
     Max,
     Model,
     Q,
     TextField,
     UUIDField,
 )
+from django.db.models.query import QuerySet
+from django.utils.translation import gettext as _
 
 Completer = t.Callable[[Context, Parameter, str], t.List[CompletionItem]]
 Strings = t.Union[t.Sequence[str], t.KeysView[str], t.Generator[str, None, None]]
@@ -107,7 +111,7 @@ def handle(
         function that returns a configured parser and completer for a model object
         and helps reduce boilerplate.
 
-    :param model_cls: The Django model class to query.
+    :param model_or_qry: The Django model class or a queryset to filter against.
     :param lookup_field: The name of the model field to use for lookup.
     :param help_field: The name of the model field to use for help text or None if
         no help text should be provided.
@@ -130,6 +134,7 @@ def handle(
     QueryBuilder = t.Callable[["ModelObjectCompleter", Context, Parameter, str], Q]
 
     model_cls: t.Type[Model]
+    _queryset: t.Optional[QuerySet] = None
     lookup_field: str
     help_field: t.Optional[str] = None
     query: t.Callable[[Context, Parameter, str], Q]
@@ -144,6 +149,10 @@ def handle(
 
     _field: Field
 
+    @property
+    def queryset(self) -> t.Union[QuerySet, Manager[Model]]:
+        return self._queryset or self.model_cls.objects
+
     def to_str(self, obj: t.Any) -> str:
         return str(obj)
 
@@ -253,7 +262,11 @@ def uuid_query(self, context: Context, parameter: Parameter, incomplete: str) ->
             self._offset += 1
 
         if len(uuid) > 32:
-            raise ValueError(f"Too many UUID characters: {incomplete}")
+            raise ValueError(
+                _("Too many UUID characters: {incomplete}").format(
+                    incomplete=incomplete
+                )
+            )
         min_uuid = UUID(uuid + "0" * (32 - len(uuid)))
         max_uuid = UUID(uuid + "f" * (32 - len(uuid)))
         return Q(**{f"{self.lookup_field}__gte": min_uuid}) & Q(
@@ -262,15 +275,23 @@ def uuid_query(self, context: Context, parameter: Parameter, incomplete: str) ->
 
     def __init__(
         self,
-        model_cls: t.Type[Model],
+        model_or_qry: t.Union[t.Type[Model], QuerySet],
         lookup_field: t.Optional[str] = None,
         help_field: t.Optional[str] = help_field,
         query: t.Optional[QueryBuilder] = None,
         limit: t.Optional[int] = limit,
         case_insensitive: bool = case_insensitive,
         distinct: bool = distinct,
     ):
-        self.model_cls = model_cls
+        if inspect.isclass(model_or_qry) and issubclass(model_or_qry, Model):
+            self.model_cls = model_or_qry
+        elif isinstance(model_or_qry, QuerySet):  # type: ignore
+            self.model_cls = model_or_qry.model
+            self._queryset = model_or_qry
+        else:
+            raise ValueError(
+                _("ModelObjectCompleter requires a Django model class or queryset.")
+            )
         self.lookup_field = str(
             lookup_field or getattr(self.model_cls._meta.pk, "name", "id")
         )
@@ -295,7 +316,9 @@ def __init__(
                 self.query = self.float_query
             else:
                 raise ValueError(
-                    f"Unsupported lookup field class: {self._field.__class__.__name__}"
+                    _("Unsupported lookup field class: {cls}").format(
+                        cls=self._field.__class__.__name__
+                    )
                 )
 
     def __call__(
@@ -343,9 +366,7 @@ def __call__(
                 ],
                 help=getattr(obj, self.help_field, None) if self.help_field else "",
             )
-            for obj in getattr(self.model_cls, "objects")
-            .filter(completion_qry)
-            .distinct()[0 : self.limit]
+            for obj in self.queryset.filter(completion_qry).distinct()[0 : self.limit]
             if (
                 getattr(obj, self.lookup_field) is not None
                 and self.to_str(getattr(obj, self.lookup_field))

django_typer/management/__init__.py

@@ -15,6 +15,7 @@
 from django.core.management.base import OutputWrapper as BaseOutputWrapper
 from django.core.management.color import Style as ColorStyle
 from django.db.models import Model
+from django.db.models.query import QuerySet
 from django.utils.functional import Promise, classproperty
 from django.utils.translation import gettext as _
 
@@ -109,7 +110,7 @@ def __call__(self, *args: P.args, **kwargs: P.kwargs) -> R:
 
 
 def model_parser_completer(
-    model_cls: t.Type[Model],
+    model_or_qry: t.Union[t.Type[Model], QuerySet],
     lookup_field: t.Optional[str] = None,
     case_insensitive: bool = False,
     help_field: t.Optional[str] = ModelObjectCompleter.help_field,
@@ -139,7 +140,7 @@ def handle(
             ...
 
 
-    :param model_cls: the model class to use for lookup
+    :param model_or_qry: the model class or QuerySet to use for lookup
     :param lookup_field: the field to use for lookup, by default the primary key
     :param case_insensitive: whether to perform case insensitive lookups and
         completions, default: False
@@ -155,13 +156,13 @@ def handle(
     """
     return {
         "parser": ModelObjectParser(
-            model_cls,
+            model_or_qry if inspect.isclass(model_or_qry) else model_or_qry.model,  # type: ignore
             lookup_field,
             case_insensitive=case_insensitive,
             on_error=on_error,
         ),
         "shell_complete": ModelObjectCompleter(
-            model_cls,
+            model_or_qry,
             lookup_field,
             case_insensitive=case_insensitive,
             help_field=help_field,

doc/source/changelog.rst

@@ -2,6 +2,11 @@
 Change Log
 ==========
 
+v2.2.0 (26-JUL-2024)
+====================
+
+* Implemented `ModelObjectCompleter should optionally accept a QuerySet in place of a Model class. <https://github.com/bckohan/django-typer/issues/96>`_
+
 v2.1.3 (15-JUL-2024)
 ====================
 

doc/source/shell_completion.rst

@@ -392,7 +392,7 @@ Model Objects
 
 * completer: :class:`~django_typer.completers.ModelObjectCompleter`
 * parser: :class:`~django_typer.parsers.ModelObjectParser`
-* convenience: :func:`~django_typer.model_parser_completer`
+* convenience: :func:`~django_typer.management.model_parser_completer`
 
 This completer/parser pairing provides the ability to fetch a model object from one of its fields.
 Most field types are supported. Additionally any other field can be set as the help text that some
@@ -419,7 +419,7 @@ shells support. Refer to the reference documentation and the
             ModelClass,
             typer.Argument(
                 **model_parser_completer(
-                    ModelClass,
+                    ModelClass,  # may also accept a QuerySet for pre-filtering
                     'field_name',  # the field that should be matched (defaults to id)
                     help_field='other_field'  # optionally provide some additional help text
                 ),

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "django-typer"
-version = "2.1.3"
+version = "2.2.0"
 description = "Use Typer to define the CLI for your Django management commands."
 authors = ["Brian Kohan <[email protected]>"]
 license = "MIT"

tests/apps/test_app/management/commands/model_fields.py

@@ -10,6 +10,7 @@
 
 import typer
 from django.utils.translation import gettext_lazy as _
+from django.db.models import Q
 
 from django_typer.management import (
     TyperCommand,
@@ -55,7 +56,7 @@ def test(
             ),
         ] = None,
         text: Annotated[
-            t.List[ShellCompleteTester],
+            t.Optional[t.List[ShellCompleteTester]],
             typer.Option(
                 **model_parser_completer(
                     ShellCompleteTester,
@@ -66,7 +67,7 @@ def test(
             ),
         ] = None,
         itext: Annotated[
-            t.List[ShellCompleteTester],
+            t.Optional[t.List[ShellCompleteTester]],
             typer.Option(
                 **model_parser_completer(
                     ShellCompleteTester,
@@ -113,26 +114,41 @@ def test(
             ),
         ] = None,
         ip: Annotated[
-            t.List[ShellCompleteTester],
+            t.Optional[t.List[ShellCompleteTester]],
             typer.Option(
                 **model_parser_completer(ShellCompleteTester, "ip_field"),
                 help=_("Fetch objects by their IP address fields."),
             ),
         ] = None,
         email: Annotated[
-            t.List[ShellCompleteTester],
+            t.Optional[t.List[ShellCompleteTester]],
             typer.Option(
                 **model_parser_completer(ShellCompleteTester, "email_field"),
                 help=_("Fetch objects by their email fields."),
             ),
         ] = None,
         url: Annotated[
-            t.List[ShellCompleteTester],
+            t.Optional[t.List[ShellCompleteTester]],
             typer.Option(
                 **model_parser_completer(ShellCompleteTester, "url_field"),
                 help=_("Fetch objects by their url fields."),
             ),
         ] = None,
+        filtered: Annotated[
+            t.Optional[t.List[ShellCompleteTester]],
+            typer.Option(
+                **model_parser_completer(
+                    ShellCompleteTester.objects.filter(
+                        ~(
+                            Q(text_field__istartswith="a")
+                            | Q(text_field__istartswith="s")
+                        )
+                    ),
+                    "text_field",
+                ),
+                help=_("Fetch objects by their text fields."),
+            ),
+        ] = None,
     ):
         assert self.__class__ is Command
         objects = {}
@@ -166,4 +182,8 @@ def test(
             for addr in ip:
                 assert isinstance(addr, ShellCompleteTester)
             objects["ip"] = [{addr.id: addr.ip_field} for addr in ip]
+        if filtered is not None:
+            for txt in filtered:
+                assert isinstance(txt, ShellCompleteTester)
+            objects["filtered"] = [{txt.id: txt.text_field} for txt in filtered]
         return json.dumps(objects)

tests/test_parser_completers.py

@@ -468,6 +468,77 @@ def test_text_field(self):
             },
         )
 
+    def test_filtered_text_field(self):
+        result = StringIO()
+        with contextlib.redirect_stdout(result):
+            call_command("shellcompletion", "complete", "model_fields test --filtered ")
+        result = result.getvalue()
+        self.assertFalse("sockeye" in result)
+        self.assertTrue("chinook" in result)
+        self.assertFalse("steelhead" in result)
+        self.assertTrue("coho" in result)
+        self.assertFalse("atlantic" in result)
+        self.assertTrue("pink" in result)
+        self.assertTrue("chum" in result)
+
+        result = StringIO()
+        with contextlib.redirect_stdout(result):
+            call_command(
+                "shellcompletion", "complete", "model_fields test --filtered ch"
+            )
+        result = result.getvalue()
+        self.assertFalse("sockeye" in result)
+        self.assertTrue("chinook" in result)
+        self.assertFalse("steelhead" in result)
+        self.assertFalse("coho" in result)
+        self.assertFalse("atlantic" in result)
+        self.assertFalse("pink" in result)
+        self.assertTrue("chum" in result)
+
+        # distinct completions by default
+        result = StringIO()
+        with contextlib.redirect_stdout(result):
+            call_command(
+                "shellcompletion",
+                "complete",
+                "model_fields test --filtered coho --filtered chinook --filtered ",
+            )
+        result = result.getvalue()
+        self.assertFalse("sockeye" in result)
+        self.assertFalse("chinook" in result)
+        self.assertFalse("steelhead" in result)
+        self.assertFalse("coho" in result)
+        self.assertFalse("atlantic" in result)
+        self.assertTrue("pink" in result)
+        self.assertTrue("chum" in result)
+
+        self.assertEqual(
+            json.loads(
+                call_command(
+                    "model_fields",
+                    "test",
+                    "--filtered",
+                    "coho",
+                    "--filtered",
+                    "chinook",
+                )
+            ),
+            {
+                "filtered": [
+                    {
+                        str(
+                            ShellCompleteTester.objects.get(text_field="coho").pk
+                        ): "coho"
+                    },
+                    {
+                        str(
+                            ShellCompleteTester.objects.get(text_field="chinook").pk
+                        ): "chinook"
+                    },
+                ]
+            },
+        )
+
     def test_uuid_field(self):
         from uuid import UUID
 
@@ -1387,3 +1458,12 @@ def test_databases_completer(self):
         ].strip()
 
         self.assertTrue("default" in result)
+
+    def test_model_completer_argument_test(self):
+        from django_typer.completers import ModelObjectCompleter
+
+        class NotAModel:
+            pass
+
+        with self.assertRaises(ValueError):
+            ModelObjectCompleter(NotAModel, "char_field", "test")