diff --git a/autocomplete/shortcuts.py b/autocomplete/shortcuts.py index 18c1b61..c5aa2ef 100644 --- a/autocomplete/shortcuts.py +++ b/autocomplete/shortcuts.py @@ -75,6 +75,12 @@ def __init__(self, queryset, label_for_record): self.queryset = queryset self.label_for_record = label_for_record + def __iter__(self, *args, **kwargs): + return (self.map_record(r) for r in self.queryset) + + def map_record(self, record): + return {"key": record.id, "label": self.label_for_record(record)} + def __getitem__(self, key): # Handle both single index and slice objects if isinstance(key, int): @@ -84,10 +90,7 @@ def __getitem__(self, key): else: raise TypeError("Invalid argument type") - mapped = [ - {"key": record.id, "label": self.label_for_record(record)} - for record in records - ] + mapped = [self.map_record(r) for r in records] if isinstance(key, int): return mapped[0] diff --git a/tests/test_model_ac.py b/tests/test_model_ac.py index 13644e1..b44a485 100644 --- a/tests/test_model_ac.py +++ b/tests/test_model_ac.py @@ -36,7 +36,7 @@ def test_model_ac_search(): ] -def test_model_ac_search_max_results(): +def test_model_ac_search_max_results(django_assert_max_num_queries): class PersonModelAC(ModelAutocomplete): model = Person search_attrs = ["name"] @@ -47,6 +47,22 @@ class PersonModelAC(ModelAutocomplete): p3 = PersonFactory(name="John3") p4 = PersonFactory(name="Jones") - results = PersonModelAC.search_items("Joh", ContextArg(None, None)) + with django_assert_max_num_queries(1): + results = PersonModelAC.search_items("Joh", ContextArg(None, None)) # should still contain 3 results, the view is responsible for truncating assert len(results) == 3 + + +def test_num_queries(django_assert_max_num_queries): + """ + regression: the queryset iterable wrapper was repeatedly evaluating the queryset + """ + + p1 = PersonFactory(name="John1") + p2 = PersonFactory(name="John2") + p3 = PersonFactory(name="John3") + p4 = PersonFactory(name="Jones") + + with django_assert_max_num_queries(1): + results = PersonModelAC.search_items("Joh", ContextArg(None, None)) + mapped_results = PersonModelAC.map_search_results(results, [])