diff --git a/flask_whooshalchemy.py b/flask_whooshalchemy.py index 79e7492..f7bf8bc 100644 --- a/flask_whooshalchemy.py +++ b/flask_whooshalchemy.py @@ -47,14 +47,19 @@ def __init__(self, query_obj, primary_key_name, ws, model, whoosh_rank=None): def __iter__(self): ''' Reorder ORM-db results according to Whoosh relevance score. ''' + super_iter = super(_QueryProxy, self).__iter__() if self._whoosh_rank is None: - return super(_QueryProxy, self).__iter__() + # Whoosh search hasn't been run so behave as normal. + return super_iter + + # Iterate through the values and re-order by whoosh relevance. ordered_by_whoosh_rank = [] - for row in super(_QueryProxy, self).__iter__(): + for row in super_iter: heapq.heappush(ordered_by_whoosh_rank, - (self._whoosh_rank[unicode(getattr(row, self._primary_key_name))], row)) + (self._whoosh_rank[unicode(getattr(row, + self._primary_key_name))], row)) def _inner(): while ordered_by_whoosh_rank: @@ -78,7 +83,9 @@ def whoosh_search(self, query): f = self.filter(getattr(self.model.__class__, self._primary_key_name).in_(result_set)) + f._whoosh_rank = result_ranks + return f