Modify and simplify logic of date_range

Author: Matt Roeschke

pandas/core/arrays/datetimes.py

@@ -32,7 +32,7 @@
 from pandas.core import ops
 
 from pandas.tseries.frequencies import to_offset
-from pandas.tseries.offsets import Tick, Day, generate_range
+from pandas.tseries.offsets import Tick, generate_range
 
 from pandas.core.arrays import datetimelike as dtl
 
@@ -239,75 +239,58 @@ def _generate_range(cls, start, end, periods, freq, tz=None,
         start, end, _normalized = _maybe_normalize_endpoints(start, end,
                                                              normalize)
 
-        tz, inferred_tz = _infer_tz_from_endpoints(start, end, tz)
+        tz, _ = _infer_tz_from_endpoints(start, end, tz)
 
-        if hasattr(freq, 'delta') and freq != Day():
-            # sub-Day Tick
-            if inferred_tz is None and tz is not None:
-                # naive dates
-                if start is not None and start.tz is None:
-                    start = start.tz_localize(tz, ambiguous=False)
+        # If we have a Timedelta-like frequency (Tick) make sure tz
+        # is set before generating the range. For relative frequencies,
+        # generate the range with naive dates.
+        localize_args = {'tz': None}
+        if isinstance(freq, Tick):
+            localize_args = {'tz': tz, 'ambiguous': False}
+        if tz is not None:
+            # Localize the start and end arguments
+            if start is not None and start.tz is None:
+                start = start.tz_localize(**localize_args)
 
-                if end is not None and end.tz is None:
-                    end = end.tz_localize(tz, ambiguous=False)
+            if end is not None and end.tz is None:
+                end = end.tz_localize(**localize_args)
 
-            if start and end:
-                if start.tz is None and end.tz is not None:
-                    start = start.tz_localize(end.tz, ambiguous=False)
+        if start and end:
+            # Make sure start and end have the same tz
+            if start.tz is None and end.tz is not None:
+                start = start.tz_localize(**localize_args)
 
-                if end.tz is None and start.tz is not None:
-                    end = end.tz_localize(start.tz, ambiguous=False)
+            if end.tz is None and start.tz is not None:
+                end = end.tz_localize(**localize_args)
 
+        if freq is not None:
             if cls._use_cached_range(freq, _normalized, start, end):
+                # Currently always False; never hit
                 index = cls._cached_range(start, end, periods=periods,
                                           freq=freq)
             else:
                 index = _generate_regular_range(cls, start, end, periods, freq)
 
-        else:
+            # TODO: Is this ever hit?
+            if tz is not None and getattr(index, 'tz', None) is None:
+                arr = conversion.tz_localize_to_utc(
+                    ensure_int64(index.values),
+                    tz, ambiguous=ambiguous)
 
-            if tz is not None:
-                # naive dates
-                if start is not None and start.tz is not None:
-                    start = start.replace(tzinfo=None)
-
-                if end is not None and end.tz is not None:
-                    end = end.replace(tzinfo=None)
-
-            if start and end:
-                if start.tz is None and end.tz is not None:
-                    end = end.replace(tzinfo=None)
-
-                if end.tz is None and start.tz is not None:
-                    start = start.replace(tzinfo=None)
-
-            if freq is not None:
-                if cls._use_cached_range(freq, _normalized, start, end):
-                    index = cls._cached_range(start, end, periods=periods,
-                                              freq=freq)
-                else:
-                    index = _generate_regular_range(cls, start, end,
-                                                    periods, freq)
-
-                if tz is not None and getattr(index, 'tz', None) is None:
-                    arr = conversion.tz_localize_to_utc(
-                        ensure_int64(index.values),
-                        tz, ambiguous=ambiguous)
-
-                    index = cls(arr)
-
-                    # index is localized datetime64 array -> have to convert
-                    # start/end as well to compare
-                    if start is not None:
-                        start = start.tz_localize(tz).asm8
-                    if end is not None:
-                        end = end.tz_localize(tz).asm8
-            else:
-                # Create a linearly spaced date_range in local time
-                start = start.tz_localize(tz)
-                end = end.tz_localize(tz)
-                arr = np.linspace(start.value, end.value, periods)
-                index = cls._simple_new(arr.astype('M8[ns]'), freq=None, tz=tz)
+                index = cls(arr)
+
+                # index is localized datetime64 array -> have to convert
+                # start/end as well to compare
+                if start is not None:
+                    start = start.tz_localize(tz).asm8
+                if end is not None:
+                    end = end.tz_localize(tz).asm8
+        else:
+            # Create a linearly spaced date_range in local time
+            start = start.tz_localize(tz)
+            end = end.tz_localize(tz)
+            arr = np.linspace(start.value, end.value, periods)
+            index = cls._simple_new(arr.astype('M8[ns]'), freq=None, tz=tz)
 
         if not left_closed and len(index) and index[0] == start:
             index = index[1:]
@@ -1255,12 +1238,10 @@ def _generate_regular_range(cls, start, end, periods, freq):
         data = cls._simple_new(data.view(_NS_DTYPE), None, tz=tz)
     else:
         tz = None
+        # Start and end should have the same timestamp by this point
         if isinstance(start, Timestamp):
             tz = start.tz
 
-        if isinstance(end, Timestamp):
-            tz = end.tz
-
         xdr = generate_range(start=start, end=end,
                              periods=periods, offset=freq)
 

pandas/tseries/offsets.py

@@ -2348,7 +2348,7 @@ class Nano(Tick):
 
 
 def generate_range(start=None, end=None, periods=None,
-                   offset=BDay(), time_rule=None):
+                   offset=BDay()):
     """
     Generates a sequence of dates corresponding to the specified time
     offset. Similar to dateutil.rrule except uses pandas DateOffset
@@ -2358,27 +2358,20 @@ def generate_range(start=None, end=None, periods=None,
     ----------
     start : datetime (default None)
     end : datetime (default None)
-    periods : int, optional
-    time_rule : (legacy) name of DateOffset object to be used, optional
-        Corresponds with names expected by tseries.frequencies.get_offset
+    periods : int, (default None)
+    offset : DateOffset, (default BDay())
 
     Notes
     -----
     * This method is faster for generating weekdays than dateutil.rrule
     * At least two of (start, end, periods) must be specified.
     * If both start and end are specified, the returned dates will
     satisfy start <= date <= end.
-    * If both time_rule and offset are specified, time_rule supersedes offset.
 
     Returns
     -------
     dates : generator object
-
     """
-    if time_rule is not None:
-        from pandas.tseries.frequencies import get_offset
-
-        offset = get_offset(time_rule)
 
     start = to_datetime(start)
     end = to_datetime(end)