Skip to content

Commit

Permalink
Add date_part() BQL function
Browse files Browse the repository at this point in the history
  • Loading branch information
dnicolodi committed Jun 9, 2024
1 parent 504a05e commit e39b7d7
Show file tree
Hide file tree
Showing 3 changed files with 47 additions and 0 deletions.
30 changes: 30 additions & 0 deletions beanquery/query_env.py
Original file line number Diff line number Diff line change
Expand Up @@ -592,6 +592,36 @@ def date_trunc(field, x):
return None


@function([str, datetime.date], int)
def date_part(field, x):
"""Extract the specified field from a date."""
if field == 'weekday' or field == 'dow':
return x.weekday()
if field == 'isoweekday' or field == 'isodow':
return x.isoweekday()
if field == 'week':
# isocalendar() returns a named tuple only in Python >= 3.9.
return x.isocalendar()[1]
if field == 'month':
return x.month
if field == 'quarter':
return (x.month - 1) // 3 + 1
if field == 'year':
return x.year
if field == 'isoyear':
# isocalendar() returns a named tuple only in Python >= 3.9.
return x.isocalendar()[0]
if field == 'decade':
return x.year // 10
if field == 'century':
return (x.year - 1) // 100 + 1
if field == 'millennium':
return (x.year - 1) // 1000 + 1
if field == 'epoch':
return int((x - datetime.date(1970, 1, 1)).total_seconds())
return None


@function([str], relativedelta)
def interval(x):
"""Construct a relative time interval."""
Expand Down
16 changes: 16 additions & 0 deletions beanquery/query_env_test.py
Original file line number Diff line number Diff line change
Expand Up @@ -206,6 +206,22 @@ def assertResult(self, expr, result):
columns, rows = query.run_query([], {}, f'SELECT {expr} FROM #')
self.assertEqual(rows[0][0], result)

def test_date_part(self):
self.assertResult('date_part("weekday", 2024-06-09)', 6)
self.assertResult('date_part("dow", 2024-06-09)', 6)
self.assertResult('date_part("isoweekday", 2024-06-09)', 7)
self.assertResult('date_part("isodow", 2024-06-09)', 7)
self.assertResult('date_part("week", 2024-06-09)', 23)
self.assertResult('date_part("month", 2024-06-09)', 6)
self.assertResult('date_part("quarter", 2024-06-09)', 2)
self.assertResult('date_part("year", 2024-06-09)', 2024)
self.assertResult('date_part("isoyear", 2024-06-09)', 2024)
self.assertResult('date_part("decade", 2024-06-09)', 202)
self.assertResult('date_part("century", 2024-06-09)', 21)
self.assertResult('date_part("millennium", 2024-06-09)', 3)
self.assertResult('date_part("epoch", 2024-06-09)', 1717891200)
self.assertResult('date_part("baz", 2024-06-09)', None)

def test_date_trunc(self):
self.assertResult('date_trunc("week", 2016-11-14)', datetime.date(2016, 11, 14)) # monday
self.assertResult('date_trunc("week", 2016-11-15)', datetime.date(2016, 11, 14)) # tuesday
Expand Down
1 change: 1 addition & 0 deletions pyproject.toml
Original file line number Diff line number Diff line change
Expand Up @@ -71,6 +71,7 @@ ignore = [
'PLR0912',
'PLR0913',
'PLR0915',
'PLR1714',
'PLR2004',
'PLW2901',
'RUF012',
Expand Down

0 comments on commit e39b7d7

Please sign in to comment.