Python 3 datetime.date subclass
import copy | |
import datetime | |
import pickle | |
# Each quarter corresponds to the following month and day combinations: | |
_q1 = (3, 31) | |
_q2 = (6, 30) | |
_q3 = (9, 30) | |
_q4 = (12, 31) | |
_qs = [_q1, _q2, _q3, _q4] | |
class Quarter(datetime.date): | |
def __new__(cls, year, q): | |
if q < 1 or q > 4: | |
raise ValueError("Quarter must be 1, 2, 3 or 4") | |
month, day = _qs[q-1] | |
return super().__new__(cls, year, month, day) | |
@property | |
def q(self): | |
return self.month // 3 # We want an int type | |
def __repr__(self): | |
return '{0}({1}, {2})'.format(self.__class__.__name__, self.year, self.q) | |
def __reduce__(self): | |
return (self.__class__, (self.year, self.q)) | |
# Some example magic methods | |
def __str__(self): | |
return "{year} Q{q}".format(year=self.year, q=self.q) | |
def __add__(self, other): | |
"""Add an integer number of quarters to the existing quarter.""" | |
assert isinstance(other, int) | |
value = self.year * 4 + self.q + other - 1 | |
return Quarter(value // 4, value % 4 + 1) | |
if __name__ == '__main__': | |
q1 = Quarter(2018, 1) | |
assert eval(repr(q1)) == q1 | |
# Copy | |
assert isinstance(copy.copy(q1), Quarter) | |
assert copy.copy(q1) == q1 | |
# Pickle | |
assert isinstance(pickle.loads(pickle.dumps(q1)), Quarter) | |
assert pickle.loads(pickle.dumps(q1)) == q1 | |
# Hashing | |
quarters = {q1: True} | |
assert quarters[q1] | |
# Other methods | |
assert q1 < Quarter(2018, 2) | |
assert q1 <= Quarter(2018, 1) | |
assert q1 > Quarter(2017, 4) | |
assert str(q1) == "2018 Q1" | |
assert (q1 + 3) == Quarter(2018, 4) | |
assert (q1 + 4) == Quarter(2019, 1) |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment