You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.

184 lines
6.5 KiB

8 years ago
  1. # -*- coding: utf-8 -*-
  2. # © 2016 Therp BV <http://therp.nl>
  3. # License AGPL-3.0 or later (http://www.gnu.org/licenses/agpl.html).
  4. import pytz
  5. from dateutil.rrule import rrule, rruleset
  6. from dateutil.tz import gettz
  7. from openerp import fields, models
  8. _RRULE_DATETIME_FIELDS = ['_until', '_dtstart']
  9. _RRULE_SCALAR_FIELDS = [
  10. '_wkst', '_cache', '_until', '_dtstart', '_count', '_freq', '_interval',
  11. ]
  12. _RRULE_ZERO_IS_NOT_NONE = ['_freq']
  13. class LocalRRuleSet(rruleset):
  14. """An rruleset that yields the naive utc representation of a date if the
  15. original date was timezone aware"""
  16. def __iter__(self):
  17. for date in rruleset.__iter__(self):
  18. if not date.tzinfo:
  19. yield date
  20. else:
  21. yield date.astimezone(pytz.utc).replace(tzinfo=None)
  22. class SerializableRRuleSet(list):
  23. """Getting our rule set json stringified is tricky, because we can't
  24. just inject our own json encoder. Now we pretend our set is a list,
  25. then json.dumps will try to iterate, which is why we can do our specific
  26. stuff in __iter__"""
  27. def __init__(self, *args):
  28. self._rrule = list(args)
  29. self.tz = None
  30. super(SerializableRRuleSet, self).__init__(self)
  31. def __iter__(self):
  32. for rule in self._rrule:
  33. yield dict(type='rrule', **{
  34. key[1:]:
  35. fields.Datetime.to_string(
  36. getattr(rule, key) if
  37. not self.tz or not getattr(rule, key) or
  38. not getattr(rule, key).tzinfo
  39. else getattr(rule, key).astimezone(pytz.utc)
  40. )
  41. if key in _RRULE_DATETIME_FIELDS
  42. else
  43. [] if getattr(rule, key) is None and
  44. key not in _RRULE_SCALAR_FIELDS
  45. else
  46. list(getattr(rule, key)) if key not in _RRULE_SCALAR_FIELDS
  47. else getattr(rule, key)
  48. for key in [
  49. '_byhour', '_wkst', '_bysecond', '_bymonthday',
  50. '_byweekno', '_bysetpos', '_cache', '_bymonth',
  51. '_byyearday', '_byweekday', '_byminute',
  52. '_until', '_dtstart', '_count', '_freq', '_interval',
  53. '_byeaster',
  54. ]
  55. })
  56. if self.tz:
  57. yield dict(type='tz', tz=self.tz)
  58. # TODO: implement rdate, exrule, exdate
  59. def __call__(self, default_self=None):
  60. """convert self to a proper rruleset for iteration.
  61. If we use defaults on our field, this will be called too with
  62. and empty recordset as parameter. In this case, we need self"""
  63. if isinstance(default_self, models.BaseModel):
  64. return self
  65. result = LocalRRuleSet()
  66. result._rrule = self._rrule
  67. return result
  68. def __nonzero__(self):
  69. return bool(self._rrule)
  70. def __repr__(self):
  71. return ', '.join(str(a) for a in self)
  72. def __getitem__(self, key):
  73. return rruleset.__getitem__(self(), key)
  74. def __getslice__(self, i, j):
  75. return rruleset.__getitem__(self(), slice(i, j))
  76. def __ne__(self, o):
  77. return not self.__eq__(o)
  78. def __eq__(self, o):
  79. return self.__repr__() == o.__repr__()
  80. def between(self, after, before, inc=False):
  81. return self().between(after, before, inc=inc)
  82. def after(self, dt, inc=False):
  83. return self().after(dt, inc=inc)
  84. def before(self, dt, inc=False):
  85. return self().before(dt, inc=inc)
  86. def count(self):
  87. return self().count()
  88. class FieldRRule(fields.Serialized):
  89. _slots = {
  90. 'stable_times': None,
  91. }
  92. def convert_to_cache(self, value, record, validate=True):
  93. result = SerializableRRuleSet()
  94. if not value:
  95. return result
  96. if isinstance(value, SerializableRRuleSet):
  97. if self.stable_times and not value.tz:
  98. self._add_tz(value, record.env.user.tz or 'utc')
  99. return value
  100. assert isinstance(value, list), 'An RRULE\'s content must be a list'
  101. tz = None
  102. for data in value:
  103. assert isinstance(data, dict), 'The list must contain dictionaries'
  104. assert 'type' in data, 'The dictionary must contain a type'
  105. data_type = data['type']
  106. if hasattr(self, 'convert_to_cache_parse_data_%s' % data_type):
  107. data = getattr(
  108. self, 'convert_to_cache_parse_data_%s' % data_type
  109. )(record, data)
  110. if data_type == 'rrule':
  111. result._rrule.append(rrule(**data))
  112. elif data_type == 'tz' and self.stable_times:
  113. tz = data['tz']
  114. # TODO: implement rdate, exrule, exdate
  115. else:
  116. raise ValueError('Unknown type given')
  117. if self.stable_times:
  118. self._add_tz(result, tz or record.env.user.tz or 'utc')
  119. return result
  120. def convert_to_cache_parse_data_rrule(self, record, data):
  121. """parse a data dictionary from the database"""
  122. return {
  123. key: fields.Datetime.from_string(value)
  124. if '_%s' % key in _RRULE_DATETIME_FIELDS
  125. else map(int, value)
  126. if value and '_%s' % key not in _RRULE_SCALAR_FIELDS
  127. else int(value) if value
  128. else None
  129. if not value and '_%s' % key not in _RRULE_ZERO_IS_NOT_NONE
  130. else value
  131. for key, value in data.iteritems()
  132. if key != 'type'
  133. }
  134. def to_column(self):
  135. """set our flag on the resulting column"""
  136. result = super(FieldRRule, self).to_column()
  137. result.stable_times = self.stable_times
  138. return result
  139. def _add_tz(self, value, tz):
  140. """set the timezone on an rruleset and adjust dates there"""
  141. value.tz = tz
  142. tz = gettz(tz)
  143. for rule in value._rrule:
  144. for fieldname in _RRULE_DATETIME_FIELDS:
  145. date = getattr(rule, fieldname)
  146. if not date:
  147. continue
  148. setattr(
  149. rule, fieldname,
  150. date.replace(tzinfo=pytz.utc).astimezone(tz),
  151. )
  152. rule._tzinfo = tz
  153. rule._timeset = tuple([
  154. rule._dtstart.replace(
  155. hour=time.hour,
  156. minute=time.minute,
  157. second=time.second,
  158. tzinfo=pytz.utc,
  159. ).astimezone(tz).timetz()
  160. if not time.tzinfo else time
  161. for time in rule._timeset
  162. ])