diff --git a/pandas/tseries/offsets.py b/pandas/tseries/offsets.py index 592cc551aeb77..bd95a62c3f2ed 100644 --- a/pandas/tseries/offsets.py +++ b/pandas/tseries/offsets.py @@ -966,7 +966,7 @@ class YearBegin(DateOffset, CacheableOffset): """DateOffset increments between calendar year begin dates""" def __init__(self, n=1, **kwds): - self.month = kwds.get('month', 12) + self.month = kwds.get('month', 1) if self.month < 1 or self.month > 12: raise ValueError('Month must go from 1 to 12') @@ -974,19 +974,44 @@ def __init__(self, n=1, **kwds): DateOffset.__init__(self, n=n, **kwds) def apply(self, other): + def _increment(date): + year = date.year + if date.month >= self.month: + year += 1 + return datetime(year, self.month, 1, date.hour, date.minute, + date.second, date.microsecond) + + def _decrement(date): + year = date.year + if date.month < self.month or (date.month == self.month and + date.day == 1): + year -= 1 + return datetime(year, self.month, 1, date.hour, date.minute, + date.second, date.microsecond) + + def _rollf(date): + if (date.month != self.month) or date.day > 1: + date = _increment(date) + return date + n = self.n - if other.month != 1 or other.day != 1: - other = datetime(other.year, 1, 1, - other.hour, other.minute, other.second, - other.microsecond) - if n <= 0: - n = n + 1 - other = other + relativedelta(years=n, day=1) - return other + result = other + if n > 0: + while n > 0: + result = _increment(result) + n -= 1 + elif n < 0: + while n < 0: + result = _decrement(result) + n += 1 + else: + # n == 0, roll forward + result = _rollf(result) - @classmethod - def onOffset(cls, dt): - return dt.month == 1 and dt.day == 1 + return result + + def onOffset(self, dt): + return dt.month == self.month and dt.day == 1 @property def rule_code(self): diff --git a/pandas/tseries/tests/test_offsets.py b/pandas/tseries/tests/test_offsets.py index f42e9b3877b3e..209f770da5c94 100644 --- a/pandas/tseries/tests/test_offsets.py +++ b/pandas/tseries/tests/test_offsets.py @@ -1128,6 +1128,7 @@ def test_offset(self): tests.append((YearBegin(-1), {datetime(2007, 1, 1): datetime(2006, 1, 1), + datetime(2007, 1, 15): datetime(2007, 1, 1), datetime(2008, 6, 30): datetime(2008, 1, 1), datetime(2008, 12, 31): datetime(2008, 1, 1), datetime(2006, 12, 29): datetime(2006, 1, 1), @@ -1139,6 +1140,26 @@ def test_offset(self): datetime(2008, 6, 30): datetime(2007, 1, 1), datetime(2008, 12, 31): datetime(2007, 1, 1), })) + tests.append((YearBegin(month=4), + {datetime(2007, 4, 1): datetime(2008, 4, 1), + datetime(2007, 4, 15): datetime(2008, 4, 1), + datetime(2007, 3, 1): datetime(2007, 4, 1), + datetime(2007, 12, 15): datetime(2008, 4, 1), + datetime(2012, 1, 31): datetime(2012, 4, 1), })) + + tests.append((YearBegin(0, month=4), + {datetime(2007, 4, 1): datetime(2007, 4, 1), + datetime(2007, 3, 1): datetime(2007, 4, 1), + datetime(2007, 12, 15): datetime(2008, 4, 1), + datetime(2012, 1, 31): datetime(2012, 4, 1), })) + + tests.append((YearBegin(-1, month=4), + {datetime(2007, 4, 1): datetime(2006, 4, 1), + datetime(2007, 3, 1): datetime(2006, 4, 1), + datetime(2007, 12, 15): datetime(2007, 4, 1), + datetime(2012, 1, 31): datetime(2011, 4, 1), })) + + for offset, cases in tests: for base, expected in cases.iteritems(): assertEq(offset, base, expected) diff --git a/pandas/tseries/tests/test_period.py b/pandas/tseries/tests/test_period.py index 5d477a8ca10ed..2183c24088bb6 100644 --- a/pandas/tseries/tests/test_period.py +++ b/pandas/tseries/tests/test_period.py @@ -1235,7 +1235,7 @@ def test_to_timestamp(self): self.assert_(result.index.equals(exp_index)) self.assertEquals(result.name, 'foo') - exp_index = date_range('1/1/2001', end='1/1/2009', freq='AS-DEC') + exp_index = date_range('1/1/2001', end='1/1/2009', freq='AS-JAN') result = series.to_timestamp(how='start') self.assert_(result.index.equals(exp_index)) @@ -1344,7 +1344,7 @@ def test_frame_to_time_stamp(self): self.assert_(result.index.equals(exp_index)) assert_almost_equal(result.values, df.values) - exp_index = date_range('1/1/2001', end='1/1/2009', freq='AS-DEC') + exp_index = date_range('1/1/2001', end='1/1/2009', freq='AS-JAN') result = df.to_timestamp('D', 'start') self.assert_(result.index.equals(exp_index)) @@ -1375,7 +1375,7 @@ def _get_with_delta(delta, freq='A-DEC'): self.assert_(result.columns.equals(exp_index)) assert_almost_equal(result.values, df.values) - exp_index = date_range('1/1/2001', end='1/1/2009', freq='AS-DEC') + exp_index = date_range('1/1/2001', end='1/1/2009', freq='AS-JAN') result = df.to_timestamp('D', 'start', axis=1) self.assert_(result.columns.equals(exp_index))