Python 3 datetime.date subclass
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
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