How to mock the “+” operator in python (specifical

问题:

I have working through some date mocking issues in Django, and have the final hurdle (I hope) is the following situation. I have a FakeDate class, which derives from datetime.date, which it mocks out.

The FakeDate class works as expected, however I get a problem when adding a datetime.timedelta to the FakeDate, in that it returns a genuine datetime.date, rather than the mock. This is important as elsewhere in a third party library there is an isinstance(value, datetime.date) check, which will always fail when using timedelta.

>>> import mock
>>> import datetime
>>>
>>> class FakeDate(datetime.date):
...     @classmethod
...     def today(cls):
...         return cls(1999, 12, 31)
...
>>> FakeDate.today()
FakeDate(1999, 12, 31)
>>> FakeDate(2000, 1, 1)
FakeDate(2000, 1, 1)
>>> FakeDate(1999, 12, 31) + datetime.timedelta(days=1)
datetime.date(2000, 1, 1)

I want the FakeDate + timedelta addition to return a FakeDate object rather than a datetime.date object - which I imagine involves patching the timedelta somehow - but how / where can I do this?

回答1:

Add a __add__ method to your FakeDate() class:

class FakeDate(datetime.date):
     @classmethod
     def today(cls):
         return cls(1999, 12, 31)
     def __add__(self, other):
         res = super(FakeDate, self).__add__(other)
         return type(self)(res.year, res.month, res.day)

Demo:

>>> class FakeDate(datetime.date):
...      @classmethod
...      def today(cls):
...          return cls(1999, 12, 31)
...      def __add__(self, other):
...          res = super(FakeDate, self).__add__(other)
...          return type(self)(res.year, res.month, res.day)
... 
>>> FakeDate.today() + datetime.timedelta(days=1)
FakeDate(2000, 1, 1)

Note that you can simply delegate the actual adding to the datetime.date class here; all we need to do is convert the result back to a FakeDate() instance.

回答2:

You just need to define an __add__ method in your FakeDate class -- it's the method that controls the behavior of the + operator.

import datetime

class FakeDate(datetime.date):
    @classmethod
    def today(cls):
        return cls(1999, 12, 31)

    def __add__(self, delta):
        # Create a datetime.date object so we don't need to do any calculations
        new_date = super(FakeDate, self).__add__(delta)
        # Then convert it to FakeDate.
        return FakeDate(new_date.year, new_date.month, new_date.day)

# Returns a FakeDate for 2000-01-01
FakeDate.today() + datetime.timedelta(days=1)

Note that this only handles the fakedate + timedelta case. If you want timedelta + fakedate to also return an instance of FakeDate, you'll need to define the __radd__ method as well (same code as __add__).

For more information on the __magic_methods__ associated with operators, see http://docs.python.org/2/reference/datamodel.html#emulating-numeric-types .