Merge pull request #11 from btubbs/master

support setting ret_type in __init__

1 files changed, 15 insertions, 11 deletions

diff --git a/src/croniter/croniter.py b/src/croniter/croniter.py
index 0f9d468..327326f 100644
--- a/src/croniter/croniter.py
+++ b/src/croniter/croniter.py
@@ -52,7 +52,8 @@ class croniter(object):
     bad_length = 'Exactly 5 or 6 columns has to be specified for iterator' \
                  'expression.'
 
-    def __init__(self, expr_format, start_time=None):
+    def __init__(self, expr_format, start_time=None, ret_type=float):
+        self._ret_type = ret_type
         if start_time is None:
             start_time = time()
 
@@ -137,13 +138,14 @@ class croniter(object):
                             else res)
         self.expanded = expanded
 
-    def get_next(self, ret_type=float):
-        return self._get_next(ret_type, is_prev=False)
+    def get_next(self, ret_type=None):
+        return self._get_next(ret_type or self._ret_type, is_prev=False)
 
-    def get_prev(self, ret_type=float):
-        return self._get_next(ret_type, is_prev=True)
+    def get_prev(self, ret_type=None):
+        return self._get_next(ret_type or self._ret_type, is_prev=True)
 
-    def get_current(self, ret_type=float):
+    def get_current(self, ret_type=None):
+        ret_type = ret_type or self._ret_type
         if ret_type == datetime.datetime:
             return self._timestamp_to_datetime(self.cur)
         return self.cur
@@ -186,24 +188,26 @@ class croniter(object):
         return self
     __next__ = next = get_next
 
-    def all_next(self, ret_type=float):
+    def all_next(self, ret_type=None):
         '''Generator of all consecutive dates. Can be used instead of
         implicit call to __iter__, whenever non-default
         'ret_type' has to be specified.
         '''
         while True:
-            yield self._get_next(ret_type, is_prev=False)
+            yield self._get_next(ret_type or self._ret_type, is_prev=False)
 
-    def all_prev(self, ret_type=float):
+    def all_prev(self, ret_type=None):
         '''Generator of all previous dates.'''
         while True:
-            yield self._get_next(ret_type, is_prev=True)
+            yield self._get_next(ret_type or self._ret_type, is_prev=True)
 
     iter = all_next  # alias, you can call .iter() instead of .all_next()
 
-    def _get_next(self, ret_type=float, is_prev=False):
+    def _get_next(self, ret_type=None, is_prev=False):
         expanded = self.expanded[:]
 
+        ret_type = ret_type or self._ret_type
+
         if ret_type not in (float, datetime.datetime):
             raise TypeError("Invalid ret_type, only 'float' or 'datetime' "
                             "is acceptable.")