diff --git a/dataset/freeze/format/fjson.py b/dataset/freeze/format/fjson.py index 0aca5ca..11a7714 100644 --- a/dataset/freeze/format/fjson.py +++ b/dataset/freeze/format/fjson.py @@ -27,7 +27,7 @@ class JSONSerializer(Serializer): result = result[0] if self.wrap: result = { - 'count': self.query.count, + 'count': len(result), 'results': result } meta = self.export.get('meta', {}) diff --git a/dataset/freeze/format/ftabson.py b/dataset/freeze/format/ftabson.py index 620bc44..279d051 100644 --- a/dataset/freeze/format/ftabson.py +++ b/dataset/freeze/format/ftabson.py @@ -13,7 +13,7 @@ class TabsonSerializer(JSONSerializer): d = [row.get(k) for k in keys] data.append(d) result = { - 'count': self.query.count, + 'count': len(data), 'fields': fields, 'data': data } diff --git a/dataset/persistence/table.py b/dataset/persistence/table.py index df359ec..0e2b84c 100644 --- a/dataset/persistence/table.py +++ b/dataset/persistence/table.py @@ -294,7 +294,7 @@ class Table(object): return self.table.c[order_by].asc() def find(self, _limit=None, _offset=0, _step=5000, - order_by='id', **_filter): + order_by='id', return_count=False, **_filter): """ Performs a simple search on the table. Simply pass keyword arguments as ``filter``. :: @@ -330,9 +330,12 @@ class Table(object): args = self._args_to_clause(_filter) # query total number of rows first - count_query = alias(self.table.select(whereclause=args, limit=_limit, offset=_offset), name='count_query_alias').count() + count_query = alias(self.table.select(whereclause=args, limit=_limit, offset=_offset), + name='count_query_alias').count() rp = self.database.executable.execute(count_query) total_row_count = rp.fetchone()[0] + if return_count: + return total_row_count if _limit is None: _limit = total_row_count @@ -355,12 +358,14 @@ class Table(object): offset=qoffset, order_by=order_by)) return ResultIter((self.database.executable.execute(q) for q in queries)) + def count(self, **_filter): + """ Return the count of results for the given filter set (same filter + options as with ``find()``). """ + return self.find(return_count=True, **_filter) + def __len__(self): - """ - Returns the number of rows in the table. - """ - d = next(self.database.query(self.table.count())) - return list(d.values()).pop() + """ Returns the number of rows in the table. """ + return self.count() def distinct(self, *columns, **_filter): """ diff --git a/dataset/persistence/util.py b/dataset/persistence/util.py index 5121618..da126c7 100644 --- a/dataset/persistence/util.py +++ b/dataset/persistence/util.py @@ -35,13 +35,11 @@ class ResultIter(object): if not isgenerator(result_proxies): result_proxies = iter((result_proxies, )) self.result_proxies = result_proxies - self.count = 0 self._iter = None def _next_rp(self): try: rp = next(self.result_proxies) - self.count += rp.rowcount self.keys = list(rp.keys()) self._iter = iter(rp.fetchall()) return True diff --git a/docs/api.rst b/docs/api.rst index 66ed530..89d4416 100644 --- a/docs/api.rst +++ b/docs/api.rst @@ -16,7 +16,7 @@ Table ----- .. autoclass:: dataset.Table - :members: columns, find, find_one, all, distinct, insert, insert_many, update, upsert, delete, create_column, drop_column, create_index, drop + :members: columns, find, find_one, all, count, distinct, insert, insert_many, update, upsert, delete, create_column, drop_column, create_index, drop :special-members: __len__, __iter__ diff --git a/test/test_persistence.py b/test/test_persistence.py index 3e0ea3c..efc3b87 100644 --- a/test/test_persistence.py +++ b/test/test_persistence.py @@ -221,6 +221,11 @@ class TableTestCase(unittest.TestCase): d = self.tbl.find_one(place='Atlantis') assert d is None, d + def test_count(self): + assert len(self.tbl) == 6, len(self.tbl) + l = self.tbl.count(place=TEST_CITY_1) + assert l == 3, l + def test_find(self): ds = list(self.tbl.find(place=TEST_CITY_1)) assert len(ds) == 3, ds