Merge pull request arrow-py#161 from zakhar/add_quarter_range

Add quarter range

arrow/arrow.py

@@ -177,6 +177,8 @@ def range(cls, frame, start, end=None, tz=None, limit=None):
         return the entire range, with **limit** alone to return a maximum # of results from the
         start, and with both to cap a range at a maximum # of results.
 
+        Supported frame values: year, quarter, month, week, day, hour, minute, second
+
         Recognized datetime expressions:
 
             - An :class:`Arrow <arrow.arrow.Arrow>` object.
@@ -204,7 +206,8 @@ def range(cls, frame, start, end=None, tz=None, limit=None):
 
         '''
 
-        frame_relative = cls._get_frames(frame)[1]
+        _, frame_relative, relative_steps = cls._get_frames(frame)
+
         tzinfo = cls._get_tzinfo(start.tzinfo if tz is None else tz)
 
         start = cls._get_datetime(start).replace(tzinfo=tzinfo)
@@ -218,7 +221,7 @@ def range(cls, frame, start, end=None, tz=None, limit=None):
             results.append(current)
 
             values = [getattr(current, f) for f in cls._ATTRS]
-            current = cls(*values, tzinfo=tzinfo) + relativedelta(**{frame_relative: 1})
+            current = cls(*values, tzinfo=tzinfo) + relativedelta(**{frame_relative: relative_steps})
 
         return results
 
@@ -238,6 +241,8 @@ def span_range(cls, frame, start, end, tz=None, limit=None):
         return the entire range, with **limit** alone to return a maximum # of results from the
         start, and with both to cap a range at a maximum # of results.
 
+        Supported frame values: year, quarter, month, week, day, hour, minute, second
+
         Recognized datetime expressions:
 
             - An :class:`Arrow <arrow.arrow.Arrow>` object.
@@ -467,6 +472,8 @@ def span(self, frame, count=1):
         :param frame: the timeframe.  Can be any ``datetime`` property (day, hour, minute...).
         :param count: (optional) the number of frames to span.
 
+        Supported frame values: year, quarter, month, week, day, hour, minute, second
+
         Usage::
 
             >>> arrow.utcnow()
@@ -483,9 +490,16 @@ def span(self, frame, count=1):
 
         '''
 
-        frame_absolute, frame_relative = self._get_frames(frame)
+        frame_absolute, frame_relative, relative_steps = self._get_frames(frame)
+
+        if frame_absolute == 'week':
+            attr = 'day'
+        elif frame_absolute == 'quarter':
+            attr = 'month'
+        else:
+            attr = frame_absolute
 
-        index = self._ATTRS.index('day' if frame_absolute == 'week' else frame_absolute)
+        index = self._ATTRS.index(attr)
         frames = self._ATTRS[:index + 1]
 
         values = [getattr(self, f) for f in frames]
@@ -497,8 +511,11 @@ def span(self, frame, count=1):
 
         if frame_absolute == 'week':
             floor = floor + relativedelta(days=-(self.isoweekday() - 1))
+        elif frame_absolute == 'quarter':
+            floor = floor + relativedelta(months=-((self.month - 1) % 3))
 
-        ceil = floor + relativedelta(**{frame_relative: count}) + relativedelta(microseconds=-1)
+        ceil = floor + relativedelta(
+            **{frame_relative: count * relative_steps}) + relativedelta(microseconds=-1)
 
         return floor, ceil
 
@@ -842,10 +859,12 @@ def _get_datetime(cls, expr):
     def _get_frames(cls, name):
 
         if name in cls._ATTRS:
-            return name, '{0}s'.format(name)
+            return name, '{0}s'.format(name), 1
 
         elif name in ['week', 'weeks']:
-            return 'week', 'weeks'
+            return 'week', 'weeks', 1
+        elif name in ['quarter', 'quarters']:
+            return 'quarter', 'months', 3
 
         raise AttributeError()
 

tests/arrow_tests.py

@@ -510,6 +510,16 @@ def test_year(self):
             arrow.Arrow(2016, 1, 2, 3, 4, 5),
         ])
 
+    def test_quarter(self):
+
+        result = arrow.Arrow.range('quarter', datetime(2013, 2, 3, 4, 5, 6),
+            datetime(2013, 5, 6, 7, 8, 9))
+
+        assertEqual(result, [
+            arrow.Arrow(2013, 2, 3, 4, 5, 6),
+            arrow.Arrow(2013, 5, 3, 4, 5, 6),
+        ])
+
     def test_month(self):
 
         result = arrow.Arrow.range('month', datetime(2013, 2, 3, 4, 5, 6),
@@ -646,6 +656,15 @@ def test_year(self):
             (arrow.Arrow(2016, 1, 1), arrow.Arrow(2016, 12, 31, 23, 59, 59, 999999)),
         ])
 
+    def test_quarter(self):
+
+        result = arrow.Arrow.span_range('quarter', datetime(2013, 2, 2), datetime(2013, 5, 15))
+
+        assertEqual(result, [
+            (arrow.Arrow(2013, 1, 1), arrow.Arrow(2013, 3, 31, 23, 59, 59, 999999)),
+            (arrow.Arrow(2013, 4, 1), arrow.Arrow(2013, 6, 30, 23, 59, 59, 999999)),
+        ])
+
     def test_month(self):
 
         result = arrow.Arrow.span_range('month', datetime(2013, 1, 2), datetime(2013, 4, 15))
@@ -783,13 +802,31 @@ def test_span_year(self):
         assertEqual(floor, datetime(2013, 1, 1, tzinfo=tz.tzutc()))
         assertEqual(ceil, datetime(2013, 12, 31, 23, 59, 59, 999999, tzinfo=tz.tzutc()))
 
+
+    def test_span_quarter(self):
+
+        floor, ceil = self.arrow.span('quarter')
+
+        assertEqual(floor, datetime(2013, 1, 1, tzinfo=tz.tzutc()))
+        assertEqual(ceil, datetime(2013, 3, 31, 23, 59, 59, 999999, tzinfo=tz.tzutc()))
+
+
+    def test_span_quarter_count(self):
+
+        floor, ceil = self.arrow.span('quarter', 2)
+
+        assertEqual(floor, datetime(2013, 1, 1, tzinfo=tz.tzutc()))
+        assertEqual(ceil, datetime(2013, 6, 30, 23, 59, 59, 999999, tzinfo=tz.tzutc()))
+
+
     def test_span_year_count(self):
 
         floor, ceil = self.arrow.span('year', 2)
 
         assertEqual(floor, datetime(2013, 1, 1, tzinfo=tz.tzutc()))
         assertEqual(ceil, datetime(2014, 12, 31, 23, 59, 59, 999999, tzinfo=tz.tzutc()))
 
+
     def test_span_month(self):
 
         floor, ceil = self.arrow.span('month')