root/django/trunk/contrib/dateutil/rrule.py

Revision 84, 39.5 kB (checked in by steadicat, 15 months ago)

Added loads of useful contrib.

  • Property svn:keywords set to Id
<
Line 
1"""
2Copyright (c) 2003-2005  Gustavo Niemeyer <gustavo@niemeyer.net>
3
4This module offers extensions to the standard python 2.3+
5datetime module.
6"""
7__author__ = "Gustavo Niemeyer <gustavo@niemeyer.net>"
8__license__ = "PSF License"
9
10import itertools
11import datetime
12import calendar
13import thread
14import sys
15
16__all__ = ["rrule", "rruleset", "rrulestr",
17           "YEARLY", "MONTHLY", "WEEKLY", "DAILY",
18           "HOURLY", "MINUTELY", "SECONDLY",
19           "MO", "TU", "WE", "TH", "FR", "SA", "SU"]
20
21# Every mask is 7 days longer to handle cross-year weekly periods.
22M366MASK = tuple([1]*31+[2]*29+[3]*31+[4]*30+[5]*31+[6]*30+
23                 [7]*31+[8]*31+[9]*30+[10]*31+[11]*30+[12]*31+[1]*7)
24M365MASK = list(M366MASK)
25M29, M30, M31 = range(1,30), range(1,31), range(1,32)
26MDAY366MASK = tuple(M31+M29+M31+M30+M31+M30+M31+M31+M30+M31+M30+M31+M31[:7])
27MDAY365MASK = list(MDAY366MASK)
28M29, M30, M31 = range(-29,0), range(-30,0), range(-31,0)
29NMDAY366MASK = tuple(M31+M29+M31+M30+M31+M30+M31+M31+M30+M31+M30+M31+M31[:7])
30NMDAY365MASK = list(NMDAY366MASK)
31M366RANGE = (0,31,60,91,121,152,182,213,244,274,305,335,366)
32M365RANGE = (0,31,59,90,120,151,181,212,243,273,304,334,365)
33WDAYMASK = [0,1,2,3,4,5,6]*55
34del M29, M30, M31, M365MASK[59], MDAY365MASK[59], NMDAY365MASK[31]
35MDAY365MASK = tuple(MDAY365MASK)
36M365MASK = tuple(M365MASK)
37
38(YEARLY,
39 MONTHLY,
40 WEEKLY,
41 DAILY,
42 HOURLY,
43 MINUTELY,
44 SECONDLY) = range(7)
45
46# Imported on demand.
47easter = None
48parser = None
49
50class weekday(object):
51    __slots__ = ["weekday", "n"]
52
53    def __init__(self, weekday, n=None):
54        if n == 0:
55            raise ValueError, "Can't create weekday with n == 0"
56        self.weekday = weekday
57        self.n = n
58
59    def __call__(self, n):
60        if n == self.n:
61            return self
62        else:
63            return self.__class__(self.weekday, n)
64
65    def __eq__(self, other):
66        try:
67            if self.weekday != other.weekday or self.n != other.n:
68                return False
69        except AttributeError:
70            return False
71        return True
72
73    def __repr__(self):
74        s = ("MO", "TU", "WE", "TH", "FR", "SA", "SU")[self.weekday]
75        if not self.n:
76            return s
77        else:
78            return "%s(%+d)" % (s, self.n)
79
80MO, TU, WE, TH, FR, SA, SU = weekdays = tuple([weekday(x) for x in range(7)])
81
82class rrulebase:
83    def __init__(self, cache=False):
84        if cache:
85            self._cache = []
86            self._cache_lock = thread.allocate_lock()
87            self._cache_gen  = self._iter()
88            self._cache_complete = False
89        else:
90            self._cache = None
91            self._cache_complete = False
92        self._len = None
93
94    def __iter__(self):
95        if self._cache_complete:
96            return iter(self._cache)
97        elif self._cache is None:
98            return self._iter()
99        else:
100            return self._iter_cached()
101
102    def _iter_cached(self):
103        i = 0
104        gen = self._cache_gen
105        cache = self._cache
106        acquire = self._cache_lock.acquire
107        release = self._cache_lock.release
108        while gen:
109            if i == len(cache):
110                acquire()
111                if self._cache_complete:
112                    break
113                try:
114                    for j in range(10):
115                        cache.append(gen.next())
116                except StopIteration:
117                    self._cache_gen = gen = None
118                    self._cache_complete = True
119                    break
120                release()
121            yield cache[i]
122            i += 1
123        while i < self._len:
124            yield cache[i]
125            i += 1
126
127    def __getitem__(self, item):
128        if self._cache_complete:
129            return self._cache[item]
130        elif isinstance(item, slice):
131            if item.step and item.step < 0:
132                return list(iter(self))[item]
133            else:
134                return list(itertools.islice(self,
135                                             item.start or 0,
136                                             item.stop or sys.maxint,
137                                             item.step or 1))
138        elif item >= 0:
139            gen = iter(self)
140            try:
141                for i in range(item+1):
142                    res = gen.next()
143            except StopIteration:
144                raise IndexError
145            return res
146        else:
147            return list(iter(self))[item]
148
149    def __contains__(self, item):
150        if self._cache_complete:
151            return item in self._cache
152        else:
153            for i in self:
154                if i == item:
155                    return True
156                elif i > item:
157                    return False
158        return False
159
160    # __len__() introduces a large performance penality.
161    def count(self):
162        if self._len is None:
163            for x in self: pass
164        return self._len
165
166    def before(self, dt, inc=False):
167        if self._cache_complete:
168            gen = self._cache
169        else:
170            gen = self
171        last = None
172        if inc:
173            for i in gen:
174                if i > dt:
175                    break
176                last = i
177        else:
178            for i in gen:
179                if i >= dt:
180                    break
181                last = i
182        return last
183
184    def after(self, dt, inc=False):
185        if self._cache_complete:
186            gen = self._cache
187        else:
188            gen = self
189        if inc:
190            for i in gen:
191                if i >= dt:
192                    return i
193        else:
194            for i in gen:
195                if i > dt:
196                    return i
197        return None
198
199    def between(self, after, before, inc=False):
200        if self._cache_complete:
201            gen = self._cache
202        else:
203            gen = self
204        started = False
205        l = []
206        if inc:
207            for i in gen:
208                if i > before:
209                    break
210                elif not started:
211                    if i >= after:
212                        started = True
213                        l.append(i)
214                else:
215                    l.append(i)
216        else:
217            for i in gen:
218                if i >= before:
219                    break
220                elif not started:
221                    if i > after:
222                        started = True
223                        l.append(i)
224                else:
225                    l.append(i)
226        return l
227
228class rrule(rrulebase):
229    def __init__(self, freq, dtstart=None,
230                 interval=1, wkst=None, count=None, until=None, bysetpos=None,
231                 bymonth=None, bymonthday=None, byyearday=None, byeaster=None,
232                 byweekno=None, byweekday=None,
233                 byhour=None, byminute=None, bysecond=None,
234                 cache=False):
235        rrulebase.__init__(self, cache)
236        global easter
237        if not dtstart:
238            dtstart = datetime.datetime.now().replace(microsecond=0)
239        elif not isinstance(dtstart, datetime.datetime):
240            dtstart = datetime.datetime.fromordinal(dtstart.toordinal())
241        else:
242            dtstart = dtstart.replace(microsecond=0)
243        self._dtstart = dtstart
244        self._tzinfo = dtstart.tzinfo
245        self._freq = freq
246        self._interval = interval
247        self._count = count
248        if until and not isinstance(until, datetime.datetime):
249            until = datetime.datetime.fromordinal(until.toordinal())
250        self._until = until
251        if wkst is None:
252            self._wkst = calendar.firstweekday()
253        elif type(wkst) is int:
254            self._wkst = wkst
255        else:
256            self._wkst = wkst.weekday
257        if bysetpos is None:
258            self._bysetpos = None
259        elif type(bysetpos) is int:
260            if bysetpos == 0 or not (-366 <= bysetpos <= 366):
261                raise ValueError("bysetpos must be between 1 and 366, "
262                                 "or between -366 and -1")
263            self._bysetpos = (bysetpos,)
264        else:
265            self._bysetpos = tuple(bysetpos)
266            for pos in self._bysetpos:
267                if pos == 0 or not (-366 <= pos <= 366):
268                    raise ValueError("bysetpos must be between 1 and 366, "
269                                     "or between -366 and -1")
270        if not (byweekno or byyearday or bymonthday or
271                byweekday is not None or byeaster is not None):
272            if freq == YEARLY:
273                if not bymonth:
274                    bymonth = dtstart.month
275                bymonthday = dtstart.day
276            elif freq == MONTHLY:
277                bymonthday = dtstart.day
278            elif freq == WEEKLY:
279                byweekday = dtstart.weekday()
280        # bymonth
281        if not bymonth:
282            self._bymonth = None
283        elif type(bymonth) is int:
284            self._bymonth = (bymonth,)
285        else:
286            self._bymonth = tuple(bymonth)
287        # byyearday
288        if not byyearday:
289            self._byyearday = None
290        elif type(byyearday) is int:
291            self._byyearday = (byyearday,)
292        else:
293            self._byyearday = tuple(byyearday)
294        # byeaster
295        if byeaster is not None:
296            if not easter:
297                from dateutil import easter
298            if type(byeaster) is int:
299                self._byeaster = (byeaster,)
300            else:
301                self._byeaster = tuple(byeaster)
302        else:
303            self._byeaster = None
304        # bymonthay
305        if not bymonthday:
306            self._bymonthday = ()
307            self._bynmonthday = ()
308        elif type(bymonthday) is int:
309            if bymonthday < 0:
310                self._bynmonthday = (bymonthday,)
311                self._bymonthday = ()
312            else:
313                self._bymonthday = (bymonthday,)
314                self._bynmonthday = ()
315        else:
316            self._bymonthday = tuple([x for x in bymonthday if x > 0])
317            self._bynmonthday = tuple([x for x in bymonthday if x < 0])
318        # byweekno
319        if byweekno is None:
320            self._byweekno = None
321        elif type(byweekno) is int:
322            self._byweekno = (byweekno,)
323        else:
324            self._byweekno = tuple(byweekno)
325        # byweekday / bynweekday
326        if byweekday is None:
327            self._byweekday = None
328            self._bynweekday = None
329        elif type(byweekday) is int:
330            self._byweekday = (byweekday,)
331            self._bynweekday = None
332        elif hasattr(byweekday, "n"):
333            if not byweekday.n or freq > MONTHLY:
334                self._byweekday = (byweekday.weekday,)
335                self._bynweekday = None
336            else:
337                self._bynweekday = ((byweekday.weekday, byweekday.n),)
338                self._byweekday = None
339        else:
340            self._byweekday = []
341            self._bynweekday = []
342            for wday in byweekday:
343                if type(wday) is int:
344                    self._byweekday.append(wday)
345                elif not wday.n or freq > MONTHLY:
346                    self._byweekday.append(wday.weekday)
347                else:
348                    self._bynweekday.append((wday.weekday, wday.n))
349            self._byweekday = tuple(self._byweekday)
350            self._bynweekday = tuple(self._bynweekday)
351            if not self._byweekday:
352                self._byweekday = None
353            elif not self._bynweekday:
354                self._bynweekday = None
355        # byhour
356        if byhour is None:
357            if freq < HOURLY:
358                self._byhour = (dtstart.hour,)
359            else:
360                self._byhour = None
361        elif type(byhour) is int:
362            self._byhour = (byhour,)
363        else:
364            self._byhour = tuple(byhour)
365        # byminute
366        if byminute is None:
367            if freq < MINUTELY:
368                self._byminute = (dtstart.minute,)
369            else:
370                self._byminute = None
371        elif type(byminute) is int:
372            self._byminute = (byminute,)
373        else:
374            self._byminute = tuple(byminute)
375        # bysecond
376        if bysecond is None:
377            if freq < SECONDLY:
378                self._bysecond = (dtstart.second,)
379            else:
380                self._bysecond = None
381        elif type(bysecond) is int:
382            self._bysecond = (bysecond,)
383        else:
384            self._bysecond = tuple(bysecond)
385
386        if self._freq >= HOURLY:
387            self._timeset = None
388        else:
389            self._timeset = []
390            for hour in self._byhour:
391                for minute in self._byminute:
392                    for second in self._bysecond:
393                        self._timeset.append(
394                                datetime.time(hour, minute, second,
395                                                    tzinfo=self._tzinfo))
396            self._timeset.sort()
397            self._timeset = tuple(self._timeset)
398
399    def _iter(self):
400        year, month, day, hour, minute, second, weekday, yearday, _ = \
401            self._dtstart.timetuple()
402
403        # Some local variables to speed things up a bit
404        freq = self._freq
405        interval = self._interval
406        wkst = self._wkst
407        until = self._until
408        bymonth = self._bymonth
409        byweekno = self._byweekno
410        byyearday = self._byyearday
411        byweekday = self._byweekday
412        byeaster = self._byeaster
413        bymonthday = self._bymonthday
414        bynmonthday = self._bynmonthday
415        bysetpos = self._bysetpos
416        byhour = self._byhour
417        byminute = self._byminute
418        bysecond = self._bysecond
419
420        ii = _iterinfo(self)
421        ii.rebuild(year, month)
422
423        getdayset = {YEARLY:ii.ydayset,
424                     MONTHLY:ii.mdayset,
425                     WEEKLY:ii.wdayset,
426                     DAILY:ii.ddayset,
427                     HOURLY:ii.ddayset,
428                     MINUTELY:ii.ddayset,
429                     SECONDLY:ii.ddayset}[freq]
430       
431        if freq < HOURLY:
432            timeset = self._timeset
433        else:
434            gettimeset = {HOURLY:ii.htimeset,
435                          MINUTELY:ii.mtimeset,
436                          SECONDLY:ii.stimeset}[freq]
437            if ((freq >= HOURLY and
438                 self._byhour and hour not in self._byhour) or
439                (freq >= MINUTELY and
440                 self._byminute and minute not in self._byminute) or
441                (freq >= SECONDLY and
442                 self._bysecond and minute not in self._bysecond)):
443                timeset = ()
444            else:
445                timeset = gettimeset(hour, minute, second)
446
447        total = 0
448        count = self._count
449        while True:
450            # Get dayset with the right frequency
451            dayset, start, end = getdayset(year, month, day)
452
453            # Do the "hard" work ;-)
454            filtered = False
455            for i in dayset[start:end]:
456                if ((bymonth and ii.mmask[i] not in bymonth) or
457                    (byweekno and not ii.wnomask[i]) or
458                    (byweekday and ii.wdaymask[i] not in byweekday) or
459                    (ii.nwdaymask and not ii.nwdaymask[i]) or
460                    (byeaster and not ii.eastermask[i]) or
461                    ((bymonthday or bynmonthday) and
462                     ii.mdaymask[i] not in bymonthday and
463                     ii.nmdaymask[i] not in bynmonthday) or
464                    (byyearday and
465                     ((i < ii.yearlen and i+1 not in byyearday
466                                      and -ii.yearlen+i not in byyearday) or
467                      (i >= ii.yearlen and i+1-ii.yearlen not in byyearday
468                                       and -ii.nextyearlen+i-ii.yearlen
469                                           not in byyearday)))):
470                    dayset[i] = None
471                    filtered = True
472
473            # Output results
474            if bysetpos and timeset:
475                poslist = []
476                for pos in bysetpos:
477                    if pos < 0:
478                        daypos, timepos = divmod(pos, len(timeset))
479                    else:
480                        daypos, timepos = divmod(pos-1, len(timeset))
481                    try:
482                        i = [x for x in dayset[start:end]
483                                if x is not None][daypos]
484                        time = timeset[timepos]
485                    except IndexError:
486                        pass
487                    else:
488                        date = datetime.date.fromordinal(ii.yearordinal+i)
489                        res = datetime.datetime.combine(date, time)
490                        if res not in poslist:
491                            poslist.append(res)
492                poslist.sort()
493                for res in poslist:
494                    if until and res > until:
495                        self._len = total
496                        return
497                    elif res >= self._dtstart:
498                        total += 1
499                        yield res
500                        if count:
501                            count -= 1
502                            if not count:
503                                self._len = total
504                                return
505            else:
506                for i in dayset[start:end]:
507                    if i is not None:
508                        date = datetime.date.fromordinal(ii.yearordinal+i)
509                        for time in timeset:
510                            res = datetime.datetime.combine(date, time)
511                            if until and res > until:
512                                self._len = total
513                                return
514                            elif res >= self._dtstart:
515                                total += 1
516                                yield res
517                                if count:
518                                    count -= 1
519                                    if not count:
520                                        self._len = total
521                                        return
522
523            # Handle frequency and interval
524            fixday = False
525            if freq == YEARLY:
526                year += interval
527                if year > datetime.MAXYEAR:
528                    self._len = total
529                    return
530                ii.rebuild(year, month)
531            elif freq == MONTHLY:
532                month += interval
533                if month > 12:
534                    div, mod = divmod(month, 12)
535                    month = mod
536                    year += div
537                    if month == 0:
538                        month = 12
539                        year -= 1
540                    if year > datetime.MAXYEAR:
541                        self._len = total
542                        return
543                ii.rebuild(year, month)
544            elif freq == WEEKLY:
545                if wkst > weekday:
546                    day += -(weekday+1+(6-wkst))+self._interval*7
547                else:
548                    day += -(weekday-wkst)+self._interval*7
549                weekday = wkst
550                fixday = True
551            elif freq == DAILY:
552                day += interval
553                fixday = True
554            elif freq == HOURLY:
555                if filtered:
556                    # Jump to one iteration before next day
557                    hour += ((23-hour)//interval)*interval
558                while True:
559                    hour += interval
560                    div, mod = divmod(hour, 24)
561                    if div:
562                        hour = mod
563                        day += div
564                        fixday = True
565                    if not byhour or hour in byhour:
566                        break
567                timeset = gettimeset(hour, minute, second)
568            elif freq == MINUTELY:
569                if filtered:
570                    # Jump to one iteration before next day
571                    minute += ((1439-(hour*60+minute))//interval)*interval
572                while True:
573                    minute += interval
574                    div, mod = divmod(minute, 60)
575                    if div:
576                        minute = mod
577                        hour += div
578                        div, mod = divmod(hour, 24)
579                        if div:
580                            hour = mod
581                            day += div
582                            fixday = True
583                            filtered = False
584                    if ((not byhour or hour in byhour) and
585                        (not byminute or minute in byminute)):
586                        break
587                timeset = gettimeset(hour, minute, second)
588            elif freq == SECONDLY:
589                if filtered:
590                    # Jump to one iteration before next day
591                    second += (((86399-(hour*3600+minute*60+second))
592                                //interval)*interval)
593                while True:
594                    second += self._interval
595                    div, mod = divmod(second, 60)
596                    if div:
597                        second = mod
598                        minute += div
599                        div, mod = divmod(minute, 60)
600                        if div:
601                            minute = mod
602                            hour += div
603                            div, mod = divmod(hour, 24)
604                            if div:
605                                hour = mod
606                                day += div
607                                fixday = True
608                    if ((not byhour or hour in byhour) and
609                        (not byminute or minute in byminute) and
610                        (not bysecond or second in bysecond)):
611                        break
612                timeset = gettimeset(hour, minute, second)
613
614            if fixday and day > 28:
615                daysinmonth = calendar.monthrange(year, month)[1]
616                if day > daysinmonth:
617                    while day > daysinmonth:
618                        day -= daysinmonth
619                        month += 1
620                        if month == 13:
621                            month = 1
622                            year += 1
623                            if year > datetime.MAXYEAR:
624                                self._len = total
625                                return
626                        daysinmonth = calendar.monthrange(year, month)[1]
627                    ii.rebuild(year, month)
628
629class _iterinfo(object):
630    __slots__ = ["rrule", "lastyear", "lastmonth",
631                 "yearlen", "nextyearlen", "yearordinal", "yearweekday",
632                 "mmask", "mrange", "mdaymask", "nmdaymask",
633                 "wdaymask", "wnomask", "nwdaymask", "eastermask"]
634
635    def __init__(self, rrule):
636        for attr in self.__slots__:
637            setattr(self, attr, None)
638        self.rrule = rrule