Skip to content

Commit 8841520

Browse files
authored
add pipe operator on QueryBuilder (#759)
* some draft work * add docstring * add test * add to docstring * add explicit string * add pipe section
1 parent 627b60a commit 8841520

File tree

3 files changed

+137
-1
lines changed

3 files changed

+137
-1
lines changed

README.rst

+52
Original file line numberDiff line numberDiff line change
@@ -1368,6 +1368,58 @@ This produces:
13681368
13691369
DROP INDEX IF EXISTS my_index
13701370
1371+
1372+
Chaining Functions
1373+
^^^^^^^^^^^^^^^^^^
1374+
1375+
The ``QueryBuilder.pipe`` method gives a more readable alternative while chaining functions.
1376+
1377+
.. code-block:: python
1378+
1379+
# This
1380+
(
1381+
query
1382+
.pipe(func1, *args)
1383+
.pipe(func2, **kwargs)
1384+
.pipe(func3)
1385+
)
1386+
1387+
# Is equivalent to this
1388+
func3(func2(func1(query, *args), **kwargs))
1389+
1390+
Or for a more concrete example:
1391+
1392+
.. code-block:: python
1393+
1394+
from pypika import Field, Query, functions as fn
1395+
from pypika.queries import QueryBuilder
1396+
1397+
def filter_days(query: QueryBuilder, col, num_days: int) -> QueryBuilder:
1398+
if isinstance(col, str):
1399+
col = Field(col)
1400+
1401+
return query.where(col > fn.Now() - num_days)
1402+
1403+
def count_groups(query: QueryBuilder, *groups) -> QueryBuilder:
1404+
return query.groupby(*groups).select(*groups, fn.Count("*").as_("n_rows"))
1405+
1406+
base_query = Query.from_("table")
1407+
1408+
query = (
1409+
base_query
1410+
.pipe(filter_days, "date", num_days=7)
1411+
.pipe(count_groups, "col1", "col2")
1412+
)
1413+
1414+
This produces:
1415+
1416+
.. code-block:: sql
1417+
1418+
SELECT "col1","col2",COUNT(*) n_rows
1419+
FROM "table"
1420+
WHERE "date">NOW()-7
1421+
GROUP BY "col1","col2"
1422+
13711423
.. _tutorial_end:
13721424

13731425
.. _contributing_start:

pypika/queries.py

+48
Original file line numberDiff line numberDiff line change
@@ -1560,6 +1560,54 @@ def _set_sql(self, **kwargs: Any) -> str:
15601560
)
15611561
)
15621562

1563+
def pipe(self, func, *args, **kwargs):
1564+
"""Call a function on the current object and return the result.
1565+
1566+
Example usage:
1567+
1568+
.. code-block:: python
1569+
1570+
from pypika import Query, functions as fn
1571+
from pypika.queries import QueryBuilder
1572+
1573+
def rows_by_group(query: QueryBuilder, *groups) -> QueryBuilder:
1574+
return (
1575+
query
1576+
.select(*groups, fn.Count("*").as_("n_rows"))
1577+
.groupby(*groups)
1578+
)
1579+
1580+
base_query = Query.from_("table")
1581+
1582+
col1_agg = base_query.pipe(rows_by_group, "col1")
1583+
col2_agg = base_query.pipe(rows_by_group, "col2")
1584+
col1_col2_agg = base_query.pipe(rows_by_group, "col1", "col2")
1585+
1586+
Makes chaining functions together easier, especially when the functions are
1587+
defined elsewhere. For example, you could define a function that filters
1588+
rows by a date range and then group by a set of columns:
1589+
1590+
1591+
.. code-block:: python
1592+
1593+
from datetime import datetime, timedelta
1594+
1595+
from pypika import Field
1596+
1597+
def days_since(query: QueryBuilder, n_days: int) -> QueryBuilder:
1598+
return (
1599+
query
1600+
.where("date" > fn.Date(datetime.now().date() - timedelta(days=n_days)))
1601+
)
1602+
1603+
(
1604+
base_query
1605+
.pipe(days_since, n_days=7)
1606+
.pipe(rows_by_group, "col1", "col2")
1607+
)
1608+
"""
1609+
return func(self, *args, **kwargs)
1610+
15631611

15641612
class Joiner:
15651613
def __init__(

pypika/tests/test_query.py

+37-1
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,6 @@
11
import unittest
22

3-
from pypika import Case, Query, Tables, Tuple, functions
3+
from pypika import Case, Query, Tables, Tuple, functions, Field
44
from pypika.dialects import (
55
ClickHouseQuery,
66
ClickHouseQueryBuilder,
@@ -204,3 +204,39 @@ def test_query_builders_have_reference_to_correct_query_class(self):
204204

205205
with self.subTest('OracleQueryBuilder'):
206206
self.assertEqual(OracleQuery, OracleQueryBuilder.QUERY_CLS)
207+
208+
def test_pipe(self) -> None:
209+
base_query = Query.from_("test")
210+
211+
def select(query: QueryBuilder) -> QueryBuilder:
212+
return query.select("test1", "test2")
213+
214+
def count_group(query: QueryBuilder, *groups) -> QueryBuilder:
215+
return query.groupby(*groups).select(*groups, functions.Count("*"))
216+
217+
for func, args, kwargs, expected_str in [
218+
(select, [], {}, 'SELECT "test1","test2" FROM "test"'),
219+
(
220+
count_group,
221+
["test1", "test2"],
222+
{},
223+
'SELECT "test1","test2",COUNT(*) FROM "test" GROUP BY "test1","test2"',
224+
),
225+
(count_group, ["test1"], {}, 'SELECT "test1",COUNT(*) FROM "test" GROUP BY "test1"'),
226+
]:
227+
result_str = str(base_query.pipe(func, *args, **kwargs))
228+
self.assertEqual(result_str, str(func(base_query, *args, **kwargs)))
229+
self.assertEqual(result_str, expected_str)
230+
231+
def where_clause(query: QueryBuilder, num_days: int) -> QueryBuilder:
232+
return query.where(Field("date") > functions.Now() - num_days)
233+
234+
result_str = str(base_query.pipe(select).pipe(where_clause, num_days=1))
235+
self.assertEqual(
236+
result_str,
237+
str(select(where_clause(base_query, num_days=1))),
238+
)
239+
self.assertEqual(
240+
result_str,
241+
'SELECT "test1","test2" FROM "test" WHERE "date">NOW()-1',
242+
)

0 commit comments

Comments
 (0)