Support multiple starting and ending times in BusinessHour

Author: Si Wei How

pandas/tseries/offsets.py

@@ -579,9 +579,28 @@ class BusinessHourMixin(BusinessMixin):
 
     def __init__(self, start='09:00', end='17:00', offset=timedelta(0)):
         # must be validated here to equality check
-        start = liboffsets._validate_business_time(start)
+        if not isinstance(start, (tuple, list)):
+            start = (start,)
+        if not isinstance(end, (tuple, list)):
+            end = (end,)
+        start = tuple(map(liboffsets._validate_business_time, start))
+        end = tuple(map(liboffsets._validate_business_time, end))
+
+        # Validation of input
+        if len(start) != len(end):
+            raise ValueError('number of starting time and ending time must be the same')
+        num_openings = len(start)
+        openings = sorted(zip(start, end))
+        start = tuple(x for x, _ in openings)
+        end = tuple(x for _, x in openings)
+        total_secs = 0
+        for i in range(num_openings):
+            total_secs += self._get_business_hours_by_sec(start[i], end[i])
+            total_secs += self._get_business_hours_by_sec(end[i], start[(i + 1) % num_openings])
+        if total_secs != 24 * 60 * 60:
+            raise ValueError('invalid starting and ending time(s)')
+
         object.__setattr__(self, "start", start)
-        end = liboffsets._validate_business_time(end)
         object.__setattr__(self, "end", end)
         object.__setattr__(self, "_offset", offset)
 
@@ -603,14 +622,8 @@ def next_bday(self):
         else:
             return BusinessDay(n=nb_offset)
 
-    @cache_readonly
-    def _get_daytime_flag(self):
-        if self.start == self.end:
-            raise ValueError('start and end must not be the same')
-        elif self.start < self.end:
-            return True
-        else:
-            return False
+    def _get_daytime_flag(self, start, end):
+        return start < end
 
     def _next_opening_time(self, other):
         """
@@ -620,44 +633,83 @@ def _next_opening_time(self, other):
         Opening time always locates on BusinessDay.
         Otherwise, closing time may not if business hour extends over midnight.
         """
+        earliest_start = self.start[0]
+        latest_start = self.start[-1]
         if not self.next_bday.onOffset(other):
             other = other + self.next_bday
+            if self.n >= 0:
+                return datetime(other.year, other.month, other.day,
+                                earliest_start.hour, earliest_start.minute)
+            else:
+                return datetime(other.year, other.month, other.day,
+                                latest_start.hour, latest_start.minute)
         else:
-            if self.n >= 0 and self.start < other.time():
+            if self.n >= 0 and latest_start < other.time():
                 other = other + self.next_bday
-            elif self.n < 0 and other.time() < self.start:
+                return datetime(other.year, other.month, other.day,
+                                earliest_start.hour, earliest_start.minute)
+            elif self.n < 0 and other.time() < earliest_start:
                 other = other + self.next_bday
-        return datetime(other.year, other.month, other.day,
-                        self.start.hour, self.start.minute)
+                return datetime(other.year, other.month, other.day,
+                                latest_start.hour, latest_start.minute)
+        if self.n >= 0:
+            for st in self.start:
+                if other.time() <= st:
+                    return datetime(other.year, other.month, other.day,
+                                    st.hour, st.minute)
+        else:
+            for st in reversed(self.start):
+                if other.time() >= st:
+                    return datetime(other.year, other.month, other.day,
+                                    st.hour, st.minute)
 
     def _prev_opening_time(self, other):
         """
         If n is positive, return yesterday's business day opening time.
         Otherwise yesterday business day's opening time.
         """
+        earliest_start = self.start[0]
+        latest_start = self.start[-1]
         if not self.next_bday.onOffset(other):
             other = other - self.next_bday
+            if self.n < 0:
+                return datetime(other.year, other.month, other.day,
+                                earliest_start.hour, earliest_start.minute)
+            else:
+                return datetime(other.year, other.month, other.day,
+                                latest_start.hour, latest_start.minute)
         else:
-            if self.n >= 0 and other.time() < self.start:
+            if self.n < 0 and latest_start < other.time():
                 other = other - self.next_bday
-            elif self.n < 0 and other.time() > self.start:
+                return datetime(other.year, other.month, other.day,
+                                earliest_start.hour, earliest_start.minute)
+            elif self.n >= 0 and other.time() < earliest_start:
                 other = other - self.next_bday
-        return datetime(other.year, other.month, other.day,
-                        self.start.hour, self.start.minute)
+                return datetime(other.year, other.month, other.day,
+                                latest_start.hour, latest_start.minute)
+        if self.n < 0:
+            for st in self.start:
+                if other.time() <= st:
+                    return datetime(other.year, other.month, other.day,
+                                    st.hour, st.minute)
+        else:
+            for st in reversed(self.start):
+                if other.time() >= st:
+                    return datetime(other.year, other.month, other.day,
+                                    st.hour, st.minute)
 
-    @cache_readonly
-    def _get_business_hours_by_sec(self):
+    def _get_business_hours_by_sec(self, start, end):
         """
         Return business hours in a day by seconds.
         """
-        if self._get_daytime_flag:
+        if self._get_daytime_flag(start, end):
             # create dummy datetime to calculate businesshours in a day
-            dtstart = datetime(2014, 4, 1, self.start.hour, self.start.minute)
-            until = datetime(2014, 4, 1, self.end.hour, self.end.minute)
+            dtstart = datetime(2014, 4, 1, start.hour, start.minute)
+            until = datetime(2014, 4, 1, end.hour, end.minute)
             return (until - dtstart).total_seconds()
         else:
-            dtstart = datetime(2014, 4, 1, self.start.hour, self.start.minute)
-            until = datetime(2014, 4, 2, self.end.hour, self.end.minute)
+            dtstart = datetime(2014, 4, 1, start.hour, start.minute)
+            until = datetime(2014, 4, 2, end.hour, end.minute)
             return (until - dtstart).total_seconds()
 
     @apply_wraps
@@ -666,13 +718,11 @@ def rollback(self, dt):
         Roll provided date backward to next offset only if not on offset.
         """
         if not self.onOffset(dt):
-            businesshours = self._get_business_hours_by_sec
             if self.n >= 0:
-                dt = self._prev_opening_time(
-                    dt) + timedelta(seconds=businesshours)
+                dt = self._prev_opening_time(dt)
             else:
-                dt = self._next_opening_time(
-                    dt) + timedelta(seconds=businesshours)
+                dt = self._next_opening_time(dt)
+            return self._get_closing_time(dt)
         return dt
 
     @apply_wraps
@@ -687,33 +737,37 @@ def rollforward(self, dt):
                 return self._prev_opening_time(dt)
         return dt
 
+    def _get_closing_time(self, dt):
+        for i, st in enumerate(self.start):
+            if st.hour == dt.hour and st.minute == dt.minute:
+                return dt + timedelta(seconds=self._get_business_hours_by_sec(st, self.end[i]))
+        raise ValueError("dt should be a starting time")
+
     @apply_wraps
     def apply(self, other):
-        daytime = self._get_daytime_flag
-        businesshours = self._get_business_hours_by_sec
-        bhdelta = timedelta(seconds=businesshours)
-
         if isinstance(other, datetime):
             # used for detecting edge condition
             nanosecond = getattr(other, 'nanosecond', 0)
             # reset timezone and nanosecond
             # other may be a Timestamp, thus not use replace
             other = datetime(other.year, other.month, other.day,
-                             other.hour, other.minute,
-                             other.second, other.microsecond)
+                    other.hour, other.minute,
+                    other.second, other.microsecond)
             n = self.n
             if n >= 0:
-                if (other.time() == self.end or
-                        not self._onOffset(other, businesshours)):
+                if (other.time() in self.end or
+                        not self._onOffset(other)):
                     other = self._next_opening_time(other)
             else:
-                if other.time() == self.start:
+                if other.time() in self.start:
                     # adjustment to move to previous business day
                     other = other - timedelta(seconds=1)
-                if not self._onOffset(other, businesshours):
+                if not self._onOffset(other):
                     other = self._next_opening_time(other)
-                    other = other + bhdelta
+                    other = self._get_closing_time(other)
 
+            businesshours = sum(self._get_business_hours_by_sec(st, en)
+                    for st, en in zip(self.start, self.end))
             bd, r = divmod(abs(n * 60), businesshours // 60)
             if n < 0:
                 bd, r = -bd, -r
@@ -722,45 +776,50 @@ def apply(self, other):
                 skip_bd = BusinessDay(n=bd)
                 # midnight business hour may not on BusinessDay
                 if not self.next_bday.onOffset(other):
-                    remain = other - self._prev_opening_time(other)
-                    other = self._next_opening_time(other + skip_bd) + remain
+                    prev_open = self._prev_opening_time(other)
+                    remain = other - prev_open
+                    other = prev_open + skip_bd + remain
                 else:
                     other = other + skip_bd
 
             hours, minutes = divmod(r, 60)
-            result = other + timedelta(hours=hours, minutes=minutes)
+            rem = timedelta(hours=hours, minutes=minutes)
 
             # because of previous adjustment, time will be larger than start
-            if ((daytime and (result.time() < self.start or
-                              self.end < result.time())) or
-                    not daytime and (self.end < result.time() < self.start)):
-                if n >= 0:
-                    bday_edge = self._prev_opening_time(other)
-                    bday_edge = bday_edge + bhdelta
-                    # calculate remainder
-                    bday_remain = result - bday_edge
-                    result = self._next_opening_time(other)
-                    result += bday_remain
-                else:
-                    bday_edge = self._next_opening_time(other)
-                    bday_remain = result - bday_edge
-                    result = self._next_opening_time(result) + bhdelta
-                    result += bday_remain
+            if n >= 0:
+                while rem != timedelta(0):
+                    bhour_left = self._get_closing_time(self._prev_opening_time(other)) - other
+                    if bhour_left >= rem:
+                        other = other + rem
+                        rem = timedelta(0)
+                    else:
+                        rem = rem - bhour_left
+                        other = self._next_opening_time(other + bhour_left)
+            else:
+                while rem != timedelta(0):
+                    bhour_left = self._next_opening_time(other) - other
+                    if bhour_left <= rem:
+                        other = other + rem
+                        rem = timedelta(0)
+                    else:
+                        rem = rem - bhour_left
+                        other = self._get_closing_time(
+                            self._next_opening_time(other + bhour_left - timedelta(seconds=1)))
+
             # edge handling
             if n >= 0:
-                if result.time() == self.end:
-                    result = self._next_opening_time(result)
+                if other.time() in self.end:
+                    other = self._next_opening_time(other)
             else:
-                if result.time() == self.start and nanosecond == 0:
+                if other.time() in self.start and nanosecond == 0:
                     # adjustment to move to previous business day
-                    result = self._next_opening_time(
-                        result - timedelta(seconds=1)) + bhdelta
+                    other = self._get_closing_time(self._next_opening_time(
+                            other - timedelta(seconds=1)))
 
-            return result
+            return other
         else:
-            # TODO: Figure out the end of this sente
             raise ApplyTypeError(
-                'Only know how to combine business hour with ')
+                    'Only know how to combine business hour with datetime')
 
     def onOffset(self, dt):
         if self.normalize and not _is_normalized(dt):
@@ -771,10 +830,9 @@ def onOffset(self, dt):
                           dt.minute, dt.second, dt.microsecond)
         # Valid BH can be on the different BusinessDay during midnight
         # Distinguish by the time spent from previous opening time
-        businesshours = self._get_business_hours_by_sec
-        return self._onOffset(dt, businesshours)
+        return self._onOffset(dt)
 
-    def _onOffset(self, dt, businesshours):
+    def _onOffset(self, dt):
         """
         Slight speedups using calculated values.
         """
@@ -787,15 +845,19 @@ def _onOffset(self, dt, businesshours):
         else:
             op = self._next_opening_time(dt)
         span = (dt - op).total_seconds()
+        businesshours = 0
+        for i, st in enumerate(self.start):
+            if op.hour == st.hour and op.minute == st.minute:
+                businesshours = self._get_business_hours_by_sec(st, self.end[i])
         if span <= businesshours:
             return True
         else:
             return False
 
     def _repr_attrs(self):
         out = super()._repr_attrs()
-        start = self.start.strftime('%H:%M')
-        end = self.end.strftime('%H:%M')
+        start = ','.join(st.strftime('%H:%M') for st in self.start)
+        end = ','.join(en.strftime('%H:%M') for en in self.end)
         attrs = ['{prefix}={start}-{end}'.format(prefix=self._prefix,
                                                  start=start, end=end)]
         out += ': ' + ', '.join(attrs)