relativedelta.py 18 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450
  1. # -*- coding: utf-8 -*-
  2. import datetime
  3. import calendar
  4. from six import integer_types
  5. __all__ = ["relativedelta", "MO", "TU", "WE", "TH", "FR", "SA", "SU"]
  6. class weekday(object):
  7. __slots__ = ["weekday", "n"]
  8. def __init__(self, weekday, n=None):
  9. self.weekday = weekday
  10. self.n = n
  11. def __call__(self, n):
  12. if n == self.n:
  13. return self
  14. else:
  15. return self.__class__(self.weekday, n)
  16. def __eq__(self, other):
  17. try:
  18. if self.weekday != other.weekday or self.n != other.n:
  19. return False
  20. except AttributeError:
  21. return False
  22. return True
  23. def __repr__(self):
  24. s = ("MO", "TU", "WE", "TH", "FR", "SA", "SU")[self.weekday]
  25. if not self.n:
  26. return s
  27. else:
  28. return "%s(%+d)" % (s, self.n)
  29. MO, TU, WE, TH, FR, SA, SU = weekdays = tuple([weekday(x) for x in range(7)])
  30. class relativedelta(object):
  31. """
  32. The relativedelta type is based on the specification of the excellent
  33. work done by M.-A. Lemburg in his
  34. `mx.DateTime <http://www.egenix.com/files/python/mxDateTime.html>`_ extension.
  35. However, notice that this type does *NOT* implement the same algorithm as
  36. his work. Do *NOT* expect it to behave like mx.DateTime's counterpart.
  37. There are two different ways to build a relativedelta instance. The
  38. first one is passing it two date/datetime classes::
  39. relativedelta(datetime1, datetime2)
  40. The second one is passing it any number of the following keyword arguments::
  41. relativedelta(arg1=x,arg2=y,arg3=z...)
  42. year, month, day, hour, minute, second, microsecond:
  43. Absolute information (argument is singular); adding or subtracting a
  44. relativedelta with absolute information does not perform an aritmetic
  45. operation, but rather REPLACES the corresponding value in the
  46. original datetime with the value(s) in relativedelta.
  47. years, months, weeks, days, hours, minutes, seconds, microseconds:
  48. Relative information, may be negative (argument is plural); adding
  49. or subtracting a relativedelta with relative information performs
  50. the corresponding aritmetic operation on the original datetime value
  51. with the information in the relativedelta.
  52. weekday:
  53. One of the weekday instances (MO, TU, etc). These instances may
  54. receive a parameter N, specifying the Nth weekday, which could
  55. be positive or negative (like MO(+1) or MO(-2). Not specifying
  56. it is the same as specifying +1. You can also use an integer,
  57. where 0=MO.
  58. leapdays:
  59. Will add given days to the date found, if year is a leap
  60. year, and the date found is post 28 of february.
  61. yearday, nlyearday:
  62. Set the yearday or the non-leap year day (jump leap days).
  63. These are converted to day/month/leapdays information.
  64. Here is the behavior of operations with relativedelta:
  65. 1. Calculate the absolute year, using the 'year' argument, or the
  66. original datetime year, if the argument is not present.
  67. 2. Add the relative 'years' argument to the absolute year.
  68. 3. Do steps 1 and 2 for month/months.
  69. 4. Calculate the absolute day, using the 'day' argument, or the
  70. original datetime day, if the argument is not present. Then,
  71. subtract from the day until it fits in the year and month
  72. found after their operations.
  73. 5. Add the relative 'days' argument to the absolute day. Notice
  74. that the 'weeks' argument is multiplied by 7 and added to
  75. 'days'.
  76. 6. Do steps 1 and 2 for hour/hours, minute/minutes, second/seconds,
  77. microsecond/microseconds.
  78. 7. If the 'weekday' argument is present, calculate the weekday,
  79. with the given (wday, nth) tuple. wday is the index of the
  80. weekday (0-6, 0=Mon), and nth is the number of weeks to add
  81. forward or backward, depending on its signal. Notice that if
  82. the calculated date is already Monday, for example, using
  83. (0, 1) or (0, -1) won't change the day.
  84. """
  85. def __init__(self, dt1=None, dt2=None,
  86. years=0, months=0, days=0, leapdays=0, weeks=0,
  87. hours=0, minutes=0, seconds=0, microseconds=0,
  88. year=None, month=None, day=None, weekday=None,
  89. yearday=None, nlyearday=None,
  90. hour=None, minute=None, second=None, microsecond=None):
  91. if dt1 and dt2:
  92. # datetime is a subclass of date. So both must be date
  93. if not (isinstance(dt1, datetime.date) and
  94. isinstance(dt2, datetime.date)):
  95. raise TypeError("relativedelta only diffs datetime/date")
  96. # We allow two dates, or two datetimes, so we coerce them to be
  97. # of the same type
  98. if (isinstance(dt1, datetime.datetime) !=
  99. isinstance(dt2, datetime.datetime)):
  100. if not isinstance(dt1, datetime.datetime):
  101. dt1 = datetime.datetime.fromordinal(dt1.toordinal())
  102. elif not isinstance(dt2, datetime.datetime):
  103. dt2 = datetime.datetime.fromordinal(dt2.toordinal())
  104. self.years = 0
  105. self.months = 0
  106. self.days = 0
  107. self.leapdays = 0
  108. self.hours = 0
  109. self.minutes = 0
  110. self.seconds = 0
  111. self.microseconds = 0
  112. self.year = None
  113. self.month = None
  114. self.day = None
  115. self.weekday = None
  116. self.hour = None
  117. self.minute = None
  118. self.second = None
  119. self.microsecond = None
  120. self._has_time = 0
  121. months = (dt1.year*12+dt1.month)-(dt2.year*12+dt2.month)
  122. self._set_months(months)
  123. dtm = self.__radd__(dt2)
  124. if dt1 < dt2:
  125. while dt1 > dtm:
  126. months += 1
  127. self._set_months(months)
  128. dtm = self.__radd__(dt2)
  129. else:
  130. while dt1 < dtm:
  131. months -= 1
  132. self._set_months(months)
  133. dtm = self.__radd__(dt2)
  134. delta = dt1 - dtm
  135. self.seconds = delta.seconds+delta.days*86400
  136. self.microseconds = delta.microseconds
  137. else:
  138. self.years = years
  139. self.months = months
  140. self.days = days+weeks*7
  141. self.leapdays = leapdays
  142. self.hours = hours
  143. self.minutes = minutes
  144. self.seconds = seconds
  145. self.microseconds = microseconds
  146. self.year = year
  147. self.month = month
  148. self.day = day
  149. self.hour = hour
  150. self.minute = minute
  151. self.second = second
  152. self.microsecond = microsecond
  153. if isinstance(weekday, integer_types):
  154. self.weekday = weekdays[weekday]
  155. else:
  156. self.weekday = weekday
  157. yday = 0
  158. if nlyearday:
  159. yday = nlyearday
  160. elif yearday:
  161. yday = yearday
  162. if yearday > 59:
  163. self.leapdays = -1
  164. if yday:
  165. ydayidx = [31, 59, 90, 120, 151, 181, 212,
  166. 243, 273, 304, 334, 366]
  167. for idx, ydays in enumerate(ydayidx):
  168. if yday <= ydays:
  169. self.month = idx+1
  170. if idx == 0:
  171. self.day = yday
  172. else:
  173. self.day = yday-ydayidx[idx-1]
  174. break
  175. else:
  176. raise ValueError("invalid year day (%d)" % yday)
  177. self._fix()
  178. def _fix(self):
  179. if abs(self.microseconds) > 999999:
  180. s = self.microseconds//abs(self.microseconds)
  181. div, mod = divmod(self.microseconds*s, 1000000)
  182. self.microseconds = mod*s
  183. self.seconds += div*s
  184. if abs(self.seconds) > 59:
  185. s = self.seconds//abs(self.seconds)
  186. div, mod = divmod(self.seconds*s, 60)
  187. self.seconds = mod*s
  188. self.minutes += div*s
  189. if abs(self.minutes) > 59:
  190. s = self.minutes//abs(self.minutes)
  191. div, mod = divmod(self.minutes*s, 60)
  192. self.minutes = mod*s
  193. self.hours += div*s
  194. if abs(self.hours) > 23:
  195. s = self.hours//abs(self.hours)
  196. div, mod = divmod(self.hours*s, 24)
  197. self.hours = mod*s
  198. self.days += div*s
  199. if abs(self.months) > 11:
  200. s = self.months//abs(self.months)
  201. div, mod = divmod(self.months*s, 12)
  202. self.months = mod*s
  203. self.years += div*s
  204. if (self.hours or self.minutes or self.seconds or self.microseconds
  205. or self.hour is not None or self.minute is not None or
  206. self.second is not None or self.microsecond is not None):
  207. self._has_time = 1
  208. else:
  209. self._has_time = 0
  210. def _set_months(self, months):
  211. self.months = months
  212. if abs(self.months) > 11:
  213. s = self.months//abs(self.months)
  214. div, mod = divmod(self.months*s, 12)
  215. self.months = mod*s
  216. self.years = div*s
  217. else:
  218. self.years = 0
  219. def __add__(self, other):
  220. if isinstance(other, relativedelta):
  221. return relativedelta(years=other.years+self.years,
  222. months=other.months+self.months,
  223. days=other.days+self.days,
  224. hours=other.hours+self.hours,
  225. minutes=other.minutes+self.minutes,
  226. seconds=other.seconds+self.seconds,
  227. microseconds=(other.microseconds +
  228. self.microseconds),
  229. leapdays=other.leapdays or self.leapdays,
  230. year=other.year or self.year,
  231. month=other.month or self.month,
  232. day=other.day or self.day,
  233. weekday=other.weekday or self.weekday,
  234. hour=other.hour or self.hour,
  235. minute=other.minute or self.minute,
  236. second=other.second or self.second,
  237. microsecond=(other.microsecond or
  238. self.microsecond))
  239. if not isinstance(other, datetime.date):
  240. raise TypeError("unsupported type for add operation")
  241. elif self._has_time and not isinstance(other, datetime.datetime):
  242. other = datetime.datetime.fromordinal(other.toordinal())
  243. year = (self.year or other.year)+self.years
  244. month = self.month or other.month
  245. if self.months:
  246. assert 1 <= abs(self.months) <= 12
  247. month += self.months
  248. if month > 12:
  249. year += 1
  250. month -= 12
  251. elif month < 1:
  252. year -= 1
  253. month += 12
  254. day = min(calendar.monthrange(year, month)[1],
  255. self.day or other.day)
  256. repl = {"year": year, "month": month, "day": day}
  257. for attr in ["hour", "minute", "second", "microsecond"]:
  258. value = getattr(self, attr)
  259. if value is not None:
  260. repl[attr] = value
  261. days = self.days
  262. if self.leapdays and month > 2 and calendar.isleap(year):
  263. days += self.leapdays
  264. ret = (other.replace(**repl)
  265. + datetime.timedelta(days=days,
  266. hours=self.hours,
  267. minutes=self.minutes,
  268. seconds=self.seconds,
  269. microseconds=self.microseconds))
  270. if self.weekday:
  271. weekday, nth = self.weekday.weekday, self.weekday.n or 1
  272. jumpdays = (abs(nth)-1)*7
  273. if nth > 0:
  274. jumpdays += (7-ret.weekday()+weekday) % 7
  275. else:
  276. jumpdays += (ret.weekday()-weekday) % 7
  277. jumpdays *= -1
  278. ret += datetime.timedelta(days=jumpdays)
  279. return ret
  280. def __radd__(self, other):
  281. return self.__add__(other)
  282. def __rsub__(self, other):
  283. return self.__neg__().__radd__(other)
  284. def __sub__(self, other):
  285. if not isinstance(other, relativedelta):
  286. raise TypeError("unsupported type for sub operation")
  287. return relativedelta(years=self.years-other.years,
  288. months=self.months-other.months,
  289. days=self.days-other.days,
  290. hours=self.hours-other.hours,
  291. minutes=self.minutes-other.minutes,
  292. seconds=self.seconds-other.seconds,
  293. microseconds=self.microseconds-other.microseconds,
  294. leapdays=self.leapdays or other.leapdays,
  295. year=self.year or other.year,
  296. month=self.month or other.month,
  297. day=self.day or other.day,
  298. weekday=self.weekday or other.weekday,
  299. hour=self.hour or other.hour,
  300. minute=self.minute or other.minute,
  301. second=self.second or other.second,
  302. microsecond=self.microsecond or other.microsecond)
  303. def __neg__(self):
  304. return relativedelta(years=-self.years,
  305. months=-self.months,
  306. days=-self.days,
  307. hours=-self.hours,
  308. minutes=-self.minutes,
  309. seconds=-self.seconds,
  310. microseconds=-self.microseconds,
  311. leapdays=self.leapdays,
  312. year=self.year,
  313. month=self.month,
  314. day=self.day,
  315. weekday=self.weekday,
  316. hour=self.hour,
  317. minute=self.minute,
  318. second=self.second,
  319. microsecond=self.microsecond)
  320. def __bool__(self):
  321. return not (not self.years and
  322. not self.months and
  323. not self.days and
  324. not self.hours and
  325. not self.minutes and
  326. not self.seconds and
  327. not self.microseconds and
  328. not self.leapdays and
  329. self.year is None and
  330. self.month is None and
  331. self.day is None and
  332. self.weekday is None and
  333. self.hour is None and
  334. self.minute is None and
  335. self.second is None and
  336. self.microsecond is None)
  337. # Compatibility with Python 2.x
  338. __nonzero__ = __bool__
  339. def __mul__(self, other):
  340. f = float(other)
  341. return relativedelta(years=int(self.years*f),
  342. months=int(self.months*f),
  343. days=int(self.days*f),
  344. hours=int(self.hours*f),
  345. minutes=int(self.minutes*f),
  346. seconds=int(self.seconds*f),
  347. microseconds=int(self.microseconds*f),
  348. leapdays=self.leapdays,
  349. year=self.year,
  350. month=self.month,
  351. day=self.day,
  352. weekday=self.weekday,
  353. hour=self.hour,
  354. minute=self.minute,
  355. second=self.second,
  356. microsecond=self.microsecond)
  357. __rmul__ = __mul__
  358. def __eq__(self, other):
  359. if not isinstance(other, relativedelta):
  360. return False
  361. if self.weekday or other.weekday:
  362. if not self.weekday or not other.weekday:
  363. return False
  364. if self.weekday.weekday != other.weekday.weekday:
  365. return False
  366. n1, n2 = self.weekday.n, other.weekday.n
  367. if n1 != n2 and not ((not n1 or n1 == 1) and (not n2 or n2 == 1)):
  368. return False
  369. return (self.years == other.years and
  370. self.months == other.months and
  371. self.days == other.days and
  372. self.hours == other.hours and
  373. self.minutes == other.minutes and
  374. self.seconds == other.seconds and
  375. self.leapdays == other.leapdays and
  376. self.year == other.year and
  377. self.month == other.month and
  378. self.day == other.day and
  379. self.hour == other.hour and
  380. self.minute == other.minute and
  381. self.second == other.second and
  382. self.microsecond == other.microsecond)
  383. def __ne__(self, other):
  384. return not self.__eq__(other)
  385. def __div__(self, other):
  386. return self.__mul__(1/float(other))
  387. __truediv__ = __div__
  388. def __repr__(self):
  389. l = []
  390. for attr in ["years", "months", "days", "leapdays",
  391. "hours", "minutes", "seconds", "microseconds"]:
  392. value = getattr(self, attr)
  393. if value:
  394. l.append("%s=%+d" % (attr, value))
  395. for attr in ["year", "month", "day", "weekday",
  396. "hour", "minute", "second", "microsecond"]:
  397. value = getattr(self, attr)
  398. if value is not None:
  399. l.append("%s=%s" % (attr, repr(value)))
  400. return "%s(%s)" % (self.__class__.__name__, ", ".join(l))
  401. # vim:ts=4:sw=4:et