diff --git a/docs/python-api.rst b/docs/python-api.rst index 7f490ab..aba43ad 100644 --- a/docs/python-api.rst +++ b/docs/python-api.rst @@ -286,6 +286,16 @@ This method also accepts ``offset=`` and ``limit=`` arguments, for specifying an ... print(row) {'id': 1, 'age': 4, 'name': 'Cleo'} +.. _python_api_rows_count_where: + +Counting rows +------------- + +To count the number of rows that would be returned by a where filter, use ``.count_where(where, where_args)``: + + >>> db["dogs"].count_where("age > ?", [1]): + 2 + .. _python_api_pks_and_rows_where: Listing rows with their primary keys @@ -1602,7 +1612,7 @@ The ``.count`` property shows the current number of rows (``select count(*) from >>> db["Street_Tree_List"].count 189144 -This property will take advantage of :ref:`python_api_cached_table_counts` if the ``use_counts_table`` property is set on the database. You can avoid that optimization entirely by calling ``table.execute_count()`` instead of accessing the property. +This property will take advantage of :ref:`python_api_cached_table_counts` if the ``use_counts_table`` property is set on the database. You can avoid that optimization entirely by calling ``table.count_where()`` instead of accessing the property. .. _python_api_introspection_columns: diff --git a/sqlite_utils/db.py b/sqlite_utils/db.py index eb714e5..b5da058 100644 --- a/sqlite_utils/db.py +++ b/sqlite_utils/db.py @@ -682,14 +682,23 @@ class Queryable: self.db = db self.name = name + def count_where( + self, + where=None, + where_args=None, + ): + sql = "select count(*) from [{}]".format(self.name) + if where is not None: + sql += " where " + where + return self.db.execute(sql, where_args or []).fetchone()[0] + def execute_count(self): - return self.db.execute( - "select count(*) from [{}]".format(self.name) - ).fetchone()[0] + # Backwards compatibility, see https://github.com/simonw/sqlite-utils/issues/305#issuecomment-890713185 + return self.count_where() @property def count(self): - return self.execute_count() + return self.count_where() @property def rows(self): @@ -820,7 +829,7 @@ class Table(Queryable): counts = self.db.cached_counts([self.name]) if counts: return next(iter(counts.values())) - return self.execute_count() + return self.count_where() def exists(self): return self.name in self.db.table_names() @@ -1719,6 +1728,8 @@ class Table(Queryable): output_type=None, drop=False, multi=False, + where=None, + where_args=None, show_progress=False, ): if isinstance(columns, str): @@ -1726,7 +1737,12 @@ class Table(Queryable): if multi: return self._convert_multi( - columns[0], fn, drop=drop, show_progress=show_progress + columns[0], + fn, + drop=drop, + where=where, + where_args=where_args, + show_progress=show_progress, ) if output is not None: @@ -1761,7 +1777,9 @@ class Table(Queryable): self.transform(drop=columns) return self - def _convert_multi(self, column, fn, drop, show_progress): + def _convert_multi( + self, column, fn, drop, show_progress, where=None, where_args=None + ): # First we execute the function pk_to_values = {} new_column_types = {} diff --git a/tests/test_enable_counts.py b/tests/test_enable_counts.py index 7a52108..d724e80 100644 --- a/tests/test_enable_counts.py +++ b/tests/test_enable_counts.py @@ -132,7 +132,7 @@ def test_uses_counts_after_enable_counts(counts_db_path): assert db["foo"].count == 1 assert logged == [ ("select name from sqlite_master where type = 'view'", None), - ("select count(*) from [foo]", None), + ("select count(*) from [foo]", []), ] logged.clear() assert not db.use_counts_table diff --git a/tests/test_introspect.py b/tests/test_introspect.py index cc33c46..dce8afc 100644 --- a/tests/test_introspect.py +++ b/tests/test_introspect.py @@ -52,7 +52,14 @@ def test_views(fresh_db): def test_count(existing_db): - assert 3 == existing_db["foo"].count + assert existing_db["foo"].count == 3 + assert existing_db["foo"].count_where() == 3 + assert existing_db["foo"].execute_count() == 3 + + +def test_count_where(existing_db): + assert existing_db["foo"].count_where("text != ?", ["two"]) == 2 + assert existing_db["foo"].count_where("text != :t", {"t": "two"}) == 2 def test_columns(existing_db):