ENH: added strftime, working to adapt tsplot.py for IntervalIndex

Author: adamklein

pandas/core/datetools.py

@@ -233,6 +233,8 @@ def __init__(self, value=None, freq=None,
                 self.freq = (base, mult)
             elif isinstance(freq, tuple):
                 self.freq = freq
+            elif isinstance(freq, (int, long)):
+                self.freq = (freq, 1)
             else:
                 raise ValueError("Expected (timerule, mult) tuple for freq")
 
@@ -425,6 +427,14 @@ def __str__(self):
         formatted = lib.skts_ordinal_to_string(self.ordinal, base, mult)
         return ("%s" % formatted)
 
+    def strftime(self, fmt):
+        base = self.freq[0]
+        mult = self.freq[1]
+        if fmt is not None:
+            return lib.skts_strftime(self.ordinal, base, mult, fmt)
+        else:
+            return lib.skts_ordinal_to_string(self.ordinal, base, mult)
+
 def _infer_interval_group(freqstr):
     return _interval_group(_reso_interval_map[freqstr])
 
@@ -433,6 +443,8 @@ def _interval_group(freqstr):
     return base // 1000 * 1000
 
 def _get_freq_code(freqstr):
+    if isinstance(freqstr, tuple):
+        return freqstr
     base, stride = _base_and_stride(freqstr)
     code = _interval_code_map[base]
     return code, stride

pandas/core/index.py

@@ -1880,6 +1880,9 @@ def __new__(cls, data=None,
             raise ValueError("Must provide freq argument if no data is "
                              "supplied")
 
+        if isinstance(freq, (int, long)):
+            freq = datetools._reverse_interval_code_map[freq]
+
         if data is None:
             start = to_interval(start, freq)
             end = to_interval(end, freq)
@@ -1896,11 +1899,11 @@ def __new__(cls, data=None,
             else:
                 data = np.arange(start.ordinal, end.ordinal+1, dtype=np.int64)
 
-            index = data.view(cls)
-            index.name = name
-            index.freq = freq
+            subarr = data.view(cls)
+            subarr.name = name
+            subarr.freq = freq
 
-            return index
+            return subarr
 
         if not isinstance(data, np.ndarray):
             if np.isscalar(data):
@@ -2071,6 +2074,21 @@ def shift(self, n):
 
         return IntervalIndex(data=self.values + n, freq=self.freq)
 
+    def __add__(self, other):
+        if isinstance(other, (int, long)):
+            return IntervalIndex(self.values + other, self.freq)
+        return super(IntervalIndex, self).__add__(other)
+
+    def __sub__(self, other):
+        if isinstance(other, (int, long)):
+            return IntervalIndex(self.values - other, self.freq)
+        if isinstance(other, Interval):
+            if other.freq != self.freq:
+                raise ValueError("Cannot do arithmetic with "
+                                 "non-conforming intervals")
+            return IntervalIndex(self.values - other.ordinal)
+        return super(IntervalIndex, self).__sub__(other)
+
     @property
     def inferred_type(self):
         # b/c data is represented as ints make sure we can't have ambiguous

pandas/src/datetime.pxd

@@ -117,6 +117,7 @@ cdef extern from "skts.h":
 
     PyObject *interval_strftime(long value, int freq, PyObject *args)
     PyObject *interval_to_string(long value, int freq)
+    PyObject *interval_to_string2(long value, int freq, char *fmt)
 
     int get_date_info(long ordinal, int freq, date_info *dinfo) except -1
 

pandas/src/datetime.pyx

@@ -554,34 +554,34 @@ cdef class DayOffset(_Offset):
                 self.t += (4 - self.dow) * us_in_day
                 self.dow = 4
 
-cdef ndarray[int64_t] _generate_range(_Offset offset, Py_ssize_t periods):
-    """
-    Generate timestamps according to offset.
-    """
-    cdef:
-        Py_ssize_t i
-        ndarray[int64_t] dtindex
-
-    dtindex = np.empty(periods, np.int64)
-    for i in range(periods):
-        dtindex[i] = offset._ts()
-        offset.next()
-    return dtindex
-
-cdef int64_t _count_range(_Offset offset, object end):
-    """
-    Count timestamps in range according to offset up to (and including)
-    end time.
-    """
-    cdef:
-        Py_ssize_t i=0
-        _TSObject e
-
-    e = convert_to_tsobject(end)
-    while offset._ts() <= e.value:
-        i += 1
-        offset.next()
-    return i
+#cdef ndarray[int64_t] _generate_range(_Offset offset, Py_ssize_t periods):
+#    """
+#    Generate timestamps according to offset.
+#    """
+#    cdef:
+#        Py_ssize_t i
+#        ndarray[int64_t] dtindex
+
+#    dtindex = np.empty(periods, np.int64)
+#    for i in range(periods):
+#        dtindex[i] = offset._ts()
+#        offset.next()
+#    return dtindex
+
+#cdef int64_t _count_range(_Offset offset, object end):
+#    """
+#    Count timestamps in range according to offset up to (and including)
+#    end time.
+#    """
+#    cdef:
+#        Py_ssize_t i=0
+#        _TSObject e
+
+#    e = convert_to_tsobject(end)
+#    while offset._ts() <= e.value:
+#        i += 1
+#        offset.next()
+#    return i
 
 # Conversion routines
 # ------------------------------------------------------------------------------
@@ -1023,6 +1023,18 @@ cpdef int64_t skts_ordinal_to_dt64(long skts_ordinal, int base, long mult):
 def skts_ordinal_to_string(long value, int base, long mult):
     return <object>interval_to_string(remove_mult(value, mult), base)
 
+def skts_strftime(long value, int freq, long mult, object fmt):
+    cdef:
+        PyObject *ptr
+
+    value = remove_mult(value, mult)
+    ptr = interval_to_string2(value, freq, <char*>fmt)
+
+    if ptr == NULL:
+        raise ValueError("Could not create string with fmt '%s'" % fmt)
+
+    return <object>ptr
+
 # interval accessors
 
 ctypedef int (*accessor)(long ordinal, int base) except -1

pandas/src/skts.c

@@ -1284,6 +1284,16 @@ PyObject *interval_to_string(long value, int freq)
     return retval;
 }
 
+PyObject *interval_to_string2(long value, int freq, char *fmt)
+{
+    PyObject *string_arg, *retval;
+    string_arg = Py_BuildValue("(s)", fmt);
+    if (string_arg == NULL) { return NULL; }
+    retval = interval_strftime(value, freq, string_arg);
+    Py_DECREF(string_arg);
+    return retval;
+}
+
 static int _quarter_year(long ordinal, int freq, int *year, int *quarter) {
     asfreq_info af_info;
     int qtr_freq;

pandas/src/skts.h

@@ -160,6 +160,7 @@ long get_python_ordinal(long skts_ordinal, int freq);
 
 PyObject *interval_strftime(long value, int freq, PyObject *args);
 PyObject *interval_to_string(long value, int freq);
+PyObject *interval_to_string2(long value, int freq, char *fmt);
 
 int get_date_info(long ordinal, int freq, struct date_info *dinfo);