From 7f7d5433f12f7428c28d20052c2f0e97466022ce Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Nguy=E1=BB=85n=20H=E1=BB=93ng=20Qu=C3=A2n?=
 <ng.hong.quan@gmail.com>
Date: Fri, 16 Dec 2016 19:34:56 +0700
Subject: [PATCH] Can pass datetime object when filtering by time

---
 README.md                 |  5 +++--
 influxalchemy/meta.py     | 19 +++++++++++++++++++
 tests/measurement_test.py | 10 ++++++++++
 tests/query_test.py       | 37 +++++++++++++++++++++++++++++++++++++
 4 files changed, 69 insertions(+), 2 deletions(-)

diff --git a/README.md b/README.md
index 703aefe..fbc33b3 100644
--- a/README.md
+++ b/README.md
@@ -129,7 +129,8 @@ flux.query(Widgets).group_by(Widgets.tag1)
 flux.query(Widgets).filter(Widgets.time > "now() - 7d")
 
 # SELECT * FROM widgets WHERE time >= '2016-01-01' AND time <= now() - 7d;
-flux.query(Widgets).filter(Widgets.time.between("'2016-01-01'", "now() - 7d"))
+d = date(2016, 1, 1)
+flux.query(Widgets).filter(Widgets.time.between(d, "now() - 7d"))
 ```
 
-Note that the date is wrapped in single-quotes (`'`).
+Note that naive datetime object will be assumed in UTC timezone.
diff --git a/influxalchemy/meta.py b/influxalchemy/meta.py
index dc050bb..7551150 100644
--- a/influxalchemy/meta.py
+++ b/influxalchemy/meta.py
@@ -1,8 +1,22 @@
 """ InfluxDB Meta Measurement. """
 
+from datetime import date, datetime, timezone
+
 from . import operations
 
 
+def make_tz_aware(d):
+    """ Make a date/datetime object to be timezone-aware """
+    # 'date' object doesn't need to be timezone aware
+    if type(d) is date:
+        return d
+    # Already aware
+    if d.tzinfo:
+        return d
+    # With naive datetime object, assume it is UTC
+    return d.replace(tzinfo=timezone.utc)
+
+
 class MetaMeasurement(type):
     """ Meta class of Measurement. """
     def __new__(mcs, name, bases, dict_):
@@ -118,6 +132,11 @@ def __init__(self, left, op, right):
         self._op = op
         lits = [operations.LK, operations.NK, operations.AND, operations.OR]
         if self._op in lits or isinstance(left, Time):
+            # If the right-hand-side value is a date/datetime object,
+            # we convert it to RFC3339 representation
+            # and wrap inside single quote
+            if isinstance(right, date):
+                right = repr(make_tz_aware(right).isoformat())
             self._right = right
         else:
             self._right = repr(right)
diff --git a/tests/measurement_test.py b/tests/measurement_test.py
index 23e0899..ca56090 100644
--- a/tests/measurement_test.py
+++ b/tests/measurement_test.py
@@ -1,5 +1,6 @@
 """ InfluxAlchemy Measurements. """
 
+from datetime import datetime
 from influxalchemy.measurement import Measurement
 from influxalchemy.meta import Tag
 from influxalchemy.meta import TagExp
@@ -131,6 +132,15 @@ def test_time_between_excl():
         TagExp(meas.time, " < ", "now() - 7d")
 
 
+def test_time_between_dt():
+    meas = Measurement.new("fizz")
+    d = datetime(2016, 1, 1)
+    exp = meas.time.between(d, "now() - 7d")
+    assert exp == \
+        TagExp(meas.time, " >= ", d) & \
+        TagExp(meas.time, " <= ", "now() - 7d")
+
+
 def test_exp_init():
     meas = Measurement.new("fizz")
     exp = TagExp(meas.buzz, " = ", "goo")
diff --git a/tests/query_test.py b/tests/query_test.py
index b212431..a1f0ce2 100644
--- a/tests/query_test.py
+++ b/tests/query_test.py
@@ -1,5 +1,7 @@
 """ InfluxAlchemy Query Tests. """
 
+from datetime import date, datetime, timedelta, timezone
+
 import influxdb
 import mock
 from influxalchemy.client import InfluxAlchemy
@@ -34,6 +36,41 @@ def test_filter(mock_qry):
     assert repr(query) == "SELECT * FROM fizz WHERE (buzz = 'goo');"
 
 
+@mock.patch("influxdb.InfluxDBClient.query")
+def test_filter_time_naive(mock_qry):
+    mock_qry.side_effect = influxdb.exceptions.InfluxDBClientError(None)
+    db = influxdb.InfluxDBClient(database="example")
+    client = InfluxAlchemy(db)
+    meas = Measurement.new("fizz")
+    d = datetime(2016, 10, 1)
+    query = client.query(meas).filter(meas.time >= d)
+    assert repr(query) == "SELECT * FROM fizz WHERE (time >= '2016-10-01T00:00:00+00:00');"
+
+
+@mock.patch("influxdb.InfluxDBClient.query")
+def test_filter_time_date(mock_qry):
+    mock_qry.side_effect = influxdb.exceptions.InfluxDBClientError(None)
+    db = influxdb.InfluxDBClient(database="example")
+    client = InfluxAlchemy(db)
+    meas = Measurement.new("fizz")
+    d = date(2016, 10, 1)
+    query = client.query(meas).filter(meas.time >= d)
+    assert repr(query) == "SELECT * FROM fizz WHERE (time >= '2016-10-01');"
+
+
+@mock.patch("influxdb.InfluxDBClient.query")
+def test_filter_time_aware(mock_qry):
+    mock_qry.side_effect = influxdb.exceptions.InfluxDBClientError(None)
+    db = influxdb.InfluxDBClient(database="example")
+    client = InfluxAlchemy(db)
+    meas = Measurement.new("fizz")
+    tz_indochina = timezone(timedelta(hours=7))
+    d_low = datetime(2016, 9, 1, tzinfo=tz_indochina)
+    d_high = datetime(2016, 10, 2, 8)
+    query = client.query(meas).filter(meas.time.between(d_low, d_high))
+    assert repr(query) == "SELECT * FROM fizz WHERE (time >= '2016-09-01T00:00:00+07:00' AND time <= '2016-10-02T08:00:00+00:00');"
+
+
 @mock.patch("influxdb.InfluxDBClient.query")
 def test_group_by(mock_qry):
     mock_qry.side_effect = influxdb.exceptions.InfluxDBClientError(None)