diff --git a/maya.py b/maya.py index 2fd0a6d..e92febf 100755 --- a/maya.py +++ b/maya.py @@ -22,6 +22,20 @@ from tzlocal import get_localzone _EPOCH_START = (1970, 1, 1) + +def validate_type_mayadt(func): + """ + Decorator to validate all the arguments to function + are of type `MayaDT` + """ + def inner(*args, **kwargs): + for arg in args + tuple(kwargs.values()): + if not isinstance(arg, MayaDT): + raise ValueError("Operation allowed only on object of type '{}'".format(MayaDT.__name__)) + return func(*args, **kwargs) + return inner + + class MayaDT(object): """The Maya Datetime object.""" @@ -36,6 +50,35 @@ class MayaDT(object): """Return's the datetime's format""" return format(self.datetime(), *args, **kwargs) + @validate_type_mayadt + def __sub__(self, maya_dt): + return MayaDT(self._epoch - maya_dt._epoch) + + @validate_type_mayadt + def __eq__(self, maya_dt): + return self._epoch == maya_dt._epoch + + @validate_type_mayadt + def __ne__(self, maya_dt): + return not self.__eq__(maya_dt) + + @validate_type_mayadt + def __lt__(self, maya_dt): + return self._epoch < maya_dt._epoch + + @validate_type_mayadt + def __le__(self, maya_dt): + return self.__lt__(maya_dt) or self.__eq__(maya_dt) + + @validate_type_mayadt + def __gt__(self, maya_dt): + return self._epoch > maya_dt._epoch + + @validate_type_mayadt + def __ge__(self, maya_dt): + return self.__gt__(maya_dt) or self.__eq__(maya_dt) + + # Timezone Crap # ------------- diff --git a/test_maya.py b/test_maya.py index 3fa6e41..6ad88e4 100755 --- a/test_maya.py +++ b/test_maya.py @@ -1,5 +1,6 @@ import pytest from datetime import datetime +import copy import maya @@ -94,3 +95,27 @@ def test_datetime_to_timezone(): dt = maya.when('2016-01-01').datetime(to_timezone='US/Eastern') assert dt.tzinfo.zone == 'US/Eastern' + +def test_comparison_operations(): + now = maya.now() + now_copy = copy.deepcopy(now) + tomorrow = maya.when('tomorrow') + + assert (now == now_copy) is True + assert (now == tomorrow) is False + + assert (now != now_copy) is False + assert (now != tomorrow) is True + + assert (now < now_copy) is False + assert (now < tomorrow) is True + + assert (now <= now_copy) is True + assert (now <= tomorrow) is True + + assert (now > now_copy) is False + assert (now > tomorrow) is False + + assert (now >= now_copy) is True + assert (now >= tomorrow) is False +