1##############################################################################
2# Copyright 2009, Gerhard Weis
3# All rights reserved.
4#
5# Redistribution and use in source and binary forms, with or without
6# modification, are permitted provided that the following conditions are met:
7#
8#  * Redistributions of source code must retain the above copyright notice,
9#    this list of conditions and the following disclaimer.
10#  * Redistributions in binary form must reproduce the above copyright notice,
11#    this list of conditions and the following disclaimer in the documentation
12#    and/or other materials provided with the distribution.
13#  * Neither the name of the authors nor the names of its contributors
14#    may be used to endorse or promote products derived from this software
15#    without specific prior written permission.
16#
17# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
18# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
19# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
20# ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS BE
21# LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
22# CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
23# SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
24# INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
25# CONTRACT, STRICT LIABILITY, OR TORT
26##############################################################################
27'''
28This module defines a Duration class.
29
30The class Duration allows to define durations in years and months and can be
31used as limited replacement for timedelta objects.
32'''
33from datetime import timedelta
34from decimal import Decimal, ROUND_FLOOR
35
36
37def fquotmod(val, low, high):
38    '''
39    A divmod function with boundaries.
40
41    '''
42    # assumes that all the maths is done with Decimals.
43    # divmod for Decimal uses truncate instead of floor as builtin
44    # divmod, so we have to do it manually here.
45    a, b = val - low, high - low
46    div = (a / b).to_integral(ROUND_FLOOR)
47    mod = a - div * b
48    # if we were not usig Decimal, it would look like this.
49    # div, mod = divmod(val - low, high - low)
50    mod += low
51    return int(div), mod
52
53
54def max_days_in_month(year, month):
55    '''
56    Determines the number of days of a specific month in a specific year.
57    '''
58    if month in (1, 3, 5, 7, 8, 10, 12):
59        return 31
60    if month in (4, 6, 9, 11):
61        return 30
62    if ((year % 400) == 0) or ((year % 100) != 0) and ((year % 4) == 0):
63        return 29
64    return 28
65
66
67class Duration(object):
68    '''
69    A class which represents a duration.
70
71    The difference to datetime.timedelta is, that this class handles also
72    differences given in years and months.
73    A Duration treats differences given in year, months separately from all
74    other components.
75
76    A Duration can be used almost like any timedelta object, however there
77    are some restrictions:
78      * It is not really possible to compare Durations, because it is unclear,
79        whether a duration of 1 year is bigger than 365 days or not.
80      * Equality is only tested between the two (year, month vs. timedelta)
81        basic components.
82
83    A Duration can also be converted into a datetime object, but this requires
84    a start date or an end date.
85
86    The algorithm to add a duration to a date is defined at
87    http://www.w3.org/TR/xmlschema-2/#adding-durations-to-dateTimes
88    '''
89
90    def __init__(self, days=0, seconds=0, microseconds=0, milliseconds=0,
91                 minutes=0, hours=0, weeks=0, months=0, years=0):
92        '''
93        Initialise this Duration instance with the given parameters.
94        '''
95        if not isinstance(months, Decimal):
96            months = Decimal(str(months))
97        if not isinstance(years, Decimal):
98            years = Decimal(str(years))
99        self.months = months
100        self.years = years
101        self.tdelta = timedelta(days, seconds, microseconds, milliseconds,
102                                minutes, hours, weeks)
103
104    def __getstate__(self):
105        return self.__dict__
106
107    def __setstate__(self, state):
108        self.__dict__.update(state)
109
110    def __getattr__(self, name):
111        '''
112        Provide direct access to attributes of included timedelta instance.
113        '''
114        return getattr(self.tdelta, name)
115
116    def __str__(self):
117        '''
118        Return a string representation of this duration similar to timedelta.
119        '''
120        params = []
121        if self.years:
122            params.append('%d years' % self.years)
123        if self.months:
124            fmt = "%d months"
125            if self.months <= 1:
126                fmt = "%d month"
127            params.append(fmt % self.months)
128        params.append(str(self.tdelta))
129        return ', '.join(params)
130
131    def __repr__(self):
132        '''
133        Return a string suitable for repr(x) calls.
134        '''
135        return "%s.%s(%d, %d, %d, years=%d, months=%d)" % (
136            self.__class__.__module__, self.__class__.__name__,
137            self.tdelta.days, self.tdelta.seconds,
138            self.tdelta.microseconds, self.years, self.months)
139
140    def __hash__(self):
141        '''
142        Return a hash of this instance so that it can be used in, for
143        example, dicts and sets.
144        '''
145        return hash((self.tdelta, self.months, self.years))
146
147    def __neg__(self):
148        """
149        A simple unary minus.
150
151        Returns a new Duration instance with all it's negated.
152        """
153        negduration = Duration(years=-self.years, months=-self.months)
154        negduration.tdelta = -self.tdelta
155        return negduration
156
157    def __add__(self, other):
158        '''
159        Durations can be added with Duration, timedelta, date and datetime
160        objects.
161        '''
162        if isinstance(other, Duration):
163            newduration = Duration(years=self.years + other.years,
164                                   months=self.months + other.months)
165            newduration.tdelta = self.tdelta + other.tdelta
166            return newduration
167        try:
168            # try anything that looks like a date or datetime
169            # 'other' has attributes year, month, day
170            # and relies on 'timedelta + other' being implemented
171            if (not(float(self.years).is_integer() and
172                    float(self.months).is_integer())):
173                raise ValueError('fractional years or months not supported'
174                                 ' for date calculations')
175            newmonth = other.month + self.months
176            carry, newmonth = fquotmod(newmonth, 1, 13)
177            newyear = other.year + self.years + carry
178            maxdays = max_days_in_month(newyear, newmonth)
179            if other.day > maxdays:
180                newday = maxdays
181            else:
182                newday = other.day
183            newdt = other.replace(year=newyear, month=newmonth, day=newday)
184            # does a timedelta + date/datetime
185            return self.tdelta + newdt
186        except AttributeError:
187            # other probably was not a date/datetime compatible object
188            pass
189        try:
190            # try if other is a timedelta
191            # relies on timedelta + timedelta supported
192            newduration = Duration(years=self.years, months=self.months)
193            newduration.tdelta = self.tdelta + other
194            return newduration
195        except AttributeError:
196            # ignore ... other probably was not a timedelta compatible object
197            pass
198        # we have tried everything .... return a NotImplemented
199        return NotImplemented
200
201    __radd__ = __add__
202
203    def __mul__(self, other):
204        if isinstance(other, int):
205            newduration = Duration(
206                years=self.years * other,
207                months=self.months * other)
208            newduration.tdelta = self.tdelta * other
209            return newduration
210        return NotImplemented
211
212    __rmul__ = __mul__
213
214    def __sub__(self, other):
215        '''
216        It is possible to subtract Duration and timedelta objects from Duration
217        objects.
218        '''
219        if isinstance(other, Duration):
220            newduration = Duration(years=self.years - other.years,
221                                   months=self.months - other.months)
222            newduration.tdelta = self.tdelta - other.tdelta
223            return newduration
224        try:
225            # do maths with our timedelta object ....
226            newduration = Duration(years=self.years, months=self.months)
227            newduration.tdelta = self.tdelta - other
228            return newduration
229        except TypeError:
230            # looks like timedelta - other is not implemented
231            pass
232        return NotImplemented
233
234    def __rsub__(self, other):
235        '''
236        It is possible to subtract Duration objecs from date, datetime and
237        timedelta objects.
238
239        TODO: there is some weird behaviour in date - timedelta ...
240              if timedelta has seconds or microseconds set, then
241              date - timedelta != date + (-timedelta)
242              for now we follow this behaviour to avoid surprises when mixing
243              timedeltas with Durations, but in case this ever changes in
244              the stdlib we can just do:
245                return -self + other
246              instead of all the current code
247        '''
248        if isinstance(other, timedelta):
249            tmpdur = Duration()
250            tmpdur.tdelta = other
251            return tmpdur - self
252        try:
253            # check if other behaves like a date/datetime object
254            # does it have year, month, day and replace?
255            if (not(float(self.years).is_integer() and
256                    float(self.months).is_integer())):
257                raise ValueError('fractional years or months not supported'
258                                 ' for date calculations')
259            newmonth = other.month - self.months
260            carry, newmonth = fquotmod(newmonth, 1, 13)
261            newyear = other.year - self.years + carry
262            maxdays = max_days_in_month(newyear, newmonth)
263            if other.day > maxdays:
264                newday = maxdays
265            else:
266                newday = other.day
267            newdt = other.replace(year=newyear, month=newmonth, day=newday)
268            return newdt - self.tdelta
269        except AttributeError:
270            # other probably was not compatible with data/datetime
271            pass
272        return NotImplemented
273
274    def __eq__(self, other):
275        '''
276        If the years, month part and the timedelta part are both equal, then
277        the two Durations are considered equal.
278        '''
279        if isinstance(other, Duration):
280            if (((self.years * 12 + self.months) ==
281                 (other.years * 12 + other.months) and
282                 self.tdelta == other.tdelta)):
283                return True
284            return False
285        # check if other con be compared against timedelta object
286        # will raise an AssertionError when optimisation is off
287        if self.years == 0 and self.months == 0:
288            return self.tdelta == other
289        return False
290
291    def __ne__(self, other):
292        '''
293        If the years, month part or the timedelta part is not equal, then
294        the two Durations are considered not equal.
295        '''
296        if isinstance(other, Duration):
297            if (((self.years * 12 + self.months) !=
298                 (other.years * 12 + other.months) or
299                 self.tdelta != other.tdelta)):
300                return True
301            return False
302        # check if other can be compared against timedelta object
303        # will raise an AssertionError when optimisation is off
304        if self.years == 0 and self.months == 0:
305            return self.tdelta != other
306        return True
307
308    def totimedelta(self, start=None, end=None):
309        '''
310        Convert this duration into a timedelta object.
311
312        This method requires a start datetime or end datetimem, but raises
313        an exception if both are given.
314        '''
315        if start is None and end is None:
316            raise ValueError("start or end required")
317        if start is not None and end is not None:
318            raise ValueError("only start or end allowed")
319        if start is not None:
320            return (start + self) - start
321        return end - (end - self)
322