jobqueue: add support for specifying next_t in datetime.datetime or datetime.time

jh0ker · jh0ker · commit 8ead72e3ef84 · 2016-12-13T23:38:13.000+01:00
diff --git a/telegram/ext/jobqueue.py b/telegram/ext/jobqueue.py
@@ -68,9 +68,12 @@ def put(self, job, next_t=None):
 
         Args:
             job (telegram.ext.Job): The ``Job`` instance representing the new job
-            next_t (Optional[int, float, datetime.timedelta]): Time in which the job
-                should be executed first. Defaults to ``job.interval``. ``int`` and ``float``
-                will be interpreted as seconds.
+            next_t (Optional[int, float, datetime.timedelta, datetime.datetime, datetime.time]):
+                Time in or at which the job should be executed first. Defaults to ``job.interval``.
+                If it is an ``int`` or a ``float``, it will be interpreted as seconds. If it is
+                a ``datetime.datetime``, it will be executed at the specified date and time. If it
+                is a ``datetime.time``, it will execute at the specified time today or, if the time
+                has already passed, tomorrow.
 
         """
         job.job_queue = self
@@ -85,8 +88,20 @@ def put(self, job, next_t=None):
             else:
                 raise ValueError("The interval argument should be of type datetime.timedelta,"
                                  " int or float")
-        elif isinstance(next_t, datetime.timedelta):
-            next_t = next_t.total_second()
+
+        elif isinstance(next_t, datetime.datetime):
+            next_t = next_t - datetime.datetime.now()
+
+        elif isinstance(next_t, datetime.time):
+            next_datetime = datetime.datetime.combine(datetime.date.today(), next_t)
+
+            if datetime.datetime.now().time() > next_t:
+                next_datetime += datetime.timedelta(days=1)
+
+            next_t = next_datetime - datetime.datetime.now()
+
+        if isinstance(next_t, datetime.timedelta):
+            next_t = next_t.total_seconds()
 
         now = time.time()
         next_t += now
@@ -227,8 +242,10 @@ class Job(object):
         callback (function): The callback function that should be executed by the Job. It should
             take two parameters ``bot`` and ``job``, where ``job`` is the ``Job`` instance. It
             can be used to terminate the job or modify its interval.
-        interval ([int, float, datetime.timedelta]): The interval in which the job will execute its
-            callback function. ``int`` and ``float`` will be interpreted as seconds.
+        interval (Optional[int, float, datetime.timedelta]): The interval in which the job will
+            execute its callback function. ``int`` and ``float`` will be interpreted as seconds.
+            If you don't set this value, you must set ``repeat=False`` and specify ``next_t`` when
+            you put the job into the job queue.
         repeat (Optional[bool]): If this job should be periodically execute its callback function
             (``True``) or only once (``False``). Defaults to ``True``
         context (Optional[object]): Additional data needed for the callback function. Can be
@@ -238,7 +255,7 @@ class Job(object):
     """
     job_queue = None
 
-    def __init__(self, callback, interval, repeat=True, context=None, days=Days.EVERY_DAY):
+    def __init__(self, callback, interval=None, repeat=True, context=None, days=Days.EVERY_DAY):
         self.callback = callback
         self.interval = interval
         self.repeat = repeat
@@ -250,10 +267,13 @@ def __init__(self, callback, interval, repeat=True, context=None, days=Days.EVER
         if not all(isinstance(day, int) for day in days):
             raise ValueError("The elements of the 'days' argument should be of type 'int'")
 
-        if not all(day >= 0 and day <= 6 for day in days):
+        if not all(0 <= day <= 6 for day in days):
             raise ValueError("The elements of the 'days' argument should be from 0 up to and "
                              "including 6")
 
+        if interval is None and repeat:
+            raise ValueError("You must either set an interval or set 'repeat' to 'False'")
+
         self.days = days
         self.name = callback.__name__
         self._remove = Event()
diff --git a/tests/test_jobqueue.py b/tests/test_jobqueue.py
@@ -62,9 +62,6 @@ def tearDown(self):
         if self.jq is not None:
             self.jq.stop()
 
-    def getSeconds(self):
-        return int(ceil(time.time()))
-
     def job1(self, bot, job):
         self.result += 1
 
@@ -79,7 +76,7 @@ def job4(self, bot, job):
         self.result += job.context
 
     def job5(self, bot, job):
-        self.job_time = self.getSeconds()
+        self.job_time = time.time()
 
     def test_basic(self):
         self.jq.put(Job(self.job1, 0.1))
@@ -181,22 +178,56 @@ def test_inUpdater(self):
 
     def test_time_unit_int(self):
         # Testing seconds in int
-        seconds_interval = 5
-        expected_time = self.getSeconds() + seconds_interval
+        seconds_interval = 2
+        expected_time = time.time() + seconds_interval
 
         self.jq.put(Job(self.job5, seconds_interval, repeat=False))
-        sleep(6)
-        self.assertEqual(self.job_time, expected_time)
+        sleep(2.5)
+        self.assertAlmostEqual(self.job_time, expected_time, delta=0.1)
 
-    def test_time_unit_dt_time(self):
+    def test_time_unit_dt_timedelta(self):
         # Testing seconds, minutes and hours as datetime.timedelta object
         # This is sufficient to test that it actually works.
-        interval = datetime.timedelta(seconds=5)
-        expected_time = self.getSeconds() + interval.total_seconds()
+        interval = datetime.timedelta(seconds=2)
+        expected_time = time.time() + interval.total_seconds()
 
         self.jq.put(Job(self.job5, interval, repeat=False))
-        sleep(6)
-        self.assertEqual(self.job_time, expected_time)
+        sleep(2.5)
+        self.assertAlmostEqual(self.job_time, expected_time, delta=0.1)
+
+    def test_time_unit_dt_datetime(self):
+        # Testing running at a specific datetime
+        delta = datetime.timedelta(seconds=2)
+        next_t = datetime.datetime.now() + delta
+        expected_time = time.time() + delta.total_seconds()
+
+        self.jq.put(Job(self.job5, repeat=False), next_t=next_t)
+        sleep(2.5)
+        self.assertAlmostEqual(self.job_time, expected_time, delta=0.1)
+
+    def test_time_unit_dt_time_today(self):
+        # Testing running at a specific time today
+        delta = 2
+        current_time = datetime.datetime.now().time()
+        next_t = datetime.time(current_time.hour, current_time.minute, current_time.second + delta,
+                               current_time.microsecond)
+        expected_time = time.time() + delta
+
+        self.jq.put(Job(self.job5, repeat=False), next_t=next_t)
+        sleep(2.5)
+        self.assertAlmostEqual(self.job_time, expected_time, delta=0.1)
+
+    def test_time_unit_dt_time_tomorrow(self):
+        # Testing running at a specific time that has passed today. Since we can't wait a day, we
+        # test if the jobs next_t has been calculated correctly
+        delta = -2
+        current_time = datetime.datetime.now().time()
+        next_t = datetime.time(current_time.hour, current_time.minute, current_time.second + delta,
+                               current_time.microsecond)
+        expected_time = time.time() + delta + 60 * 60 * 24
+
+        self.jq.put(Job(self.job5, repeat=False), next_t=next_t)
+        self.assertAlmostEqual(self.jq.queue.get(False)[0], expected_time, delta=0.1)
 
 
 if __name__ == '__main__':
