16
16
17
17
# pylint: disable=line-too-long
18
18
19
+ import inspect
19
20
import os
20
21
import pkgutil
21
22
import sys
38
39
FloatField ,
39
40
GenericIPAddressField ,
40
41
IntegerField ,
42
+ Manager ,
41
43
Max ,
42
44
Model ,
43
45
Q ,
44
46
TextField ,
45
47
UUIDField ,
46
48
)
49
+ from django .db .models .query import QuerySet
50
+ from django .utils .translation import gettext as _
47
51
48
52
Completer = t .Callable [[Context , Parameter , str ], t .List [CompletionItem ]]
49
53
Strings = t .Union [t .Sequence [str ], t .KeysView [str ], t .Generator [str , None , None ]]
@@ -107,7 +111,7 @@ def handle(
107
111
function that returns a configured parser and completer for a model object
108
112
and helps reduce boilerplate.
109
113
110
- :param model_cls : The Django model class to query .
114
+ :param model_or_qry : The Django model class or a queryset to filter against .
111
115
:param lookup_field: The name of the model field to use for lookup.
112
116
:param help_field: The name of the model field to use for help text or None if
113
117
no help text should be provided.
@@ -130,6 +134,7 @@ def handle(
130
134
QueryBuilder = t .Callable [["ModelObjectCompleter" , Context , Parameter , str ], Q ]
131
135
132
136
model_cls : t .Type [Model ]
137
+ _queryset : t .Optional [QuerySet ] = None
133
138
lookup_field : str
134
139
help_field : t .Optional [str ] = None
135
140
query : t .Callable [[Context , Parameter , str ], Q ]
@@ -144,6 +149,10 @@ def handle(
144
149
145
150
_field : Field
146
151
152
+ @property
153
+ def queryset (self ) -> t .Union [QuerySet , Manager [Model ]]:
154
+ return self ._queryset or self .model_cls .objects
155
+
147
156
def to_str (self , obj : t .Any ) -> str :
148
157
return str (obj )
149
158
@@ -253,7 +262,11 @@ def uuid_query(self, context: Context, parameter: Parameter, incomplete: str) ->
253
262
self ._offset += 1
254
263
255
264
if len (uuid ) > 32 :
256
- raise ValueError (f"Too many UUID characters: { incomplete } " )
265
+ raise ValueError (
266
+ _ ("Too many UUID characters: {incomplete}" ).format (
267
+ incomplete = incomplete
268
+ )
269
+ )
257
270
min_uuid = UUID (uuid + "0" * (32 - len (uuid )))
258
271
max_uuid = UUID (uuid + "f" * (32 - len (uuid )))
259
272
return Q (** {f"{ self .lookup_field } __gte" : min_uuid }) & Q (
@@ -262,15 +275,23 @@ def uuid_query(self, context: Context, parameter: Parameter, incomplete: str) ->
262
275
263
276
def __init__ (
264
277
self ,
265
- model_cls : t .Type [Model ],
278
+ model_or_qry : t .Union [ t . Type [Model ], QuerySet ],
266
279
lookup_field : t .Optional [str ] = None ,
267
280
help_field : t .Optional [str ] = help_field ,
268
281
query : t .Optional [QueryBuilder ] = None ,
269
282
limit : t .Optional [int ] = limit ,
270
283
case_insensitive : bool = case_insensitive ,
271
284
distinct : bool = distinct ,
272
285
):
273
- self .model_cls = model_cls
286
+ if inspect .isclass (model_or_qry ) and issubclass (model_or_qry , Model ):
287
+ self .model_cls = model_or_qry
288
+ elif isinstance (model_or_qry , QuerySet ): # type: ignore
289
+ self .model_cls = model_or_qry .model
290
+ self ._queryset = model_or_qry
291
+ else :
292
+ raise ValueError (
293
+ _ ("ModelObjectCompleter requires a Django model class or queryset." )
294
+ )
274
295
self .lookup_field = str (
275
296
lookup_field or getattr (self .model_cls ._meta .pk , "name" , "id" )
276
297
)
@@ -295,7 +316,9 @@ def __init__(
295
316
self .query = self .float_query
296
317
else :
297
318
raise ValueError (
298
- f"Unsupported lookup field class: { self ._field .__class__ .__name__ } "
319
+ _ ("Unsupported lookup field class: {cls}" ).format (
320
+ cls = self ._field .__class__ .__name__
321
+ )
299
322
)
300
323
301
324
def __call__ (
@@ -343,9 +366,7 @@ def __call__(
343
366
],
344
367
help = getattr (obj , self .help_field , None ) if self .help_field else "" ,
345
368
)
346
- for obj in getattr (self .model_cls , "objects" )
347
- .filter (completion_qry )
348
- .distinct ()[0 : self .limit ]
369
+ for obj in self .queryset .filter (completion_qry ).distinct ()[0 : self .limit ]
349
370
if (
350
371
getattr (obj , self .lookup_field ) is not None
351
372
and self .to_str (getattr (obj , self .lookup_field ))
0 commit comments