From e266d926f394110c634f2300cf58b569b14429b8 Mon Sep 17 00:00:00 2001 From: Eddie Hebert Date: Wed, 28 Nov 2012 10:53:08 -0500 Subject: [PATCH] Updates quarter utilities to fix offset. Also, includes tests of quarter utilities. Thanks to @richafrank. --- tests/test_quarter_utils.py | 49 +++++++++++++++++++++++++++++++++++++ zipline/utils/date_utils.py | 18 +++----------- 2 files changed, 53 insertions(+), 14 deletions(-) create mode 100644 tests/test_quarter_utils.py diff --git a/tests/test_quarter_utils.py b/tests/test_quarter_utils.py new file mode 100644 index 00000000..3cdfeb02 --- /dev/null +++ b/tests/test_quarter_utils.py @@ -0,0 +1,49 @@ +# +# Copyright 2012 Quantopian, Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + + +from unittest import TestCase + +import datetime +import pytz + +from zipline.utils.date_utils import ( + get_quarter, + dates_of_quarter +) + + +class DateUtilsQuarterTests(TestCase): + + def test_dates_quarter_inverses(self): + last_quarter = None + + for y in xrange(1900, 2050): + for m in xrange(1, 13): + + dt = datetime.datetime(y, m, 1, tzinfo=pytz.utc) + + q = get_quarter(dt) + if last_quarter: + self.assertGreater(q, last_quarter) + + boundaries = dates_of_quarter(q) + self.assertTrue( + boundaries[0] <= dt <= boundaries[1], + "dates_of_quarter not inverse of get_quarter {0}" + .format(dt,)) + + self.assertEqual(q, get_quarter(boundaries[0])) + self.assertEqual(q, get_quarter(boundaries[1])) diff --git a/zipline/utils/date_utils.py b/zipline/utils/date_utils.py index dea620b6..877b0a5b 100644 --- a/zipline/utils/date_utils.py +++ b/zipline/utils/date_utils.py @@ -107,23 +107,13 @@ def get_quarter(dt): convert the given datetime to an integer representing the number of calendar quarters since 0. """ - quarters = dt.year * 4 - month = dt.month - if month <= 3: - return quarters + 1 - elif month <= 6: - return quarters + 2 - elif month <= 9: - return quarters + 3 - else: - return quarters + 4 + return (dt.year - 1) * 4 + (dt.month - 1) / 3 + 1 def dates_of_quarter(quarter_num): - year = quarter_num / 4 - quarter = quarter_num % 4 - if quarter == 0: - quarter = 4 + quarter_num -= 1 + year = quarter_num / 4 + 1 + quarter = quarter_num % 4 + 1 if quarter == 1: start = datetime(year, 1, 1, 0, 0, tzinfo=pytz.utc)