forked from seatgeek/businesstime
-
Notifications
You must be signed in to change notification settings - Fork 0
/
__init__.py
154 lines (129 loc) · 5.71 KB
/
__init__.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
import datetime
__version__ = "0.2.0"
class BusinessTime(object):
"""
BusinessTime is essentially a calendar that can be queried for
business time aware timedeltas between two datetimes.
"""
def __init__(self, business_hours=None, weekends=(5,6), holidays=None):
if business_hours is None:
business_hours = (datetime.time(9), datetime.time(17))
self.business_hours = business_hours
# TODO: weekends should maybe be a generator or a callable returning True/False
self.weekends = weekends
self.holidays = holidays
if callable(self.holidays) or self.holidays is None:
self._holidaysGeneratorStart = None
self._holidaysGenerator = None
self._holidays = []
else:
self._holidays = self.holidays
# HACK: pick an arbitrary date so we can do math with datetime.time objects
arbitrary_date = datetime.datetime(2014, 1, 26)
start = datetime.datetime.combine(arbitrary_date, business_hours[0])
end = datetime.datetime.combine(arbitrary_date, business_hours[1])
self.open_hours = end - start
def isweekend(self, dt):
return dt.weekday() in self.weekends
def _ensure_holidays_span_datetime(self, dt):
if callable(self.holidays):
if self._holidaysGeneratorStart is None or dt < self._holidaysGeneratorStart:
self._holidaysGeneratorStart = dt
self._holidaysGenerator = self.holidays(dt)
while len(self._holidays) == 0 or dt > self._holidays[-1]:
self._holidays.append(next(self._holidaysGenerator))
def isholiday(self, dt):
if type(dt) == datetime.datetime:
dt = dt.date()
self._ensure_holidays_span_datetime(dt)
return dt in self._holidays
def isbusinessday(self, dt):
return not self.isweekend(dt) and not self.isholiday(dt)
def isduringbusinesshours(self, dt):
return self.isbusinessday(dt) and self.business_hours[0] <= dt.time() < self.business_hours[1]
def iterdays(self, d1, d2):
"""
Date iterator returning dates in d1 <= x < d2
"""
curr = datetime.datetime.combine(d1, datetime.time())
end = datetime.datetime.combine(d2, datetime.time())
if d1.date() == d2.date():
yield curr
return
while curr < end:
yield curr
curr = curr + datetime.timedelta(days=1)
def iterweekdays(self, d1, d2):
"""
Date iterator returning dates in d1 <= x < d2, excluding weekends
"""
for dt in self.iterdays(d1, d2):
if not self.isweekend(dt):
yield dt
def iterbusinessdays(self, d1, d2):
"""
Date iterator returning dates in d1 <= x < d2, excluding weekends and holidays
"""
assert d2 >= d1
if d1.date() == d2.date() and d2.time() < self.business_hours[0]:
return
first = True
for dt in self.iterdays(d1, d2):
if first and d1.time() > self.business_hours[1]:
first = False
continue
first = False
if not self.isweekend(dt) and not self.isholiday(dt):
yield dt
def _build_spanning_datetimes(self, d1, d2):
businessdays = list(self.iterbusinessdays(d1, d2))
if len(businessdays) == 0:
return businessdays
businessdays = [datetime.datetime.combine(d, self.business_hours[0]) for d in businessdays]
if d1 > businessdays[0]:
businessdays[0] = d1
if self.isbusinessday(d2) and d2 >= datetime.datetime.combine(d2, self.business_hours[0]):
businessdays.append(datetime.datetime.combine(d2, self.business_hours[1]))
if d2 < businessdays[-1]:
businessdays[-1] = datetime.datetime.combine(businessdays[-1], d2.time())
else:
if len(businessdays) == 1:
businessdays.append(datetime.datetime.combine(businessdays[0], self.business_hours[1]))
else:
businessdays[-1] = datetime.datetime.combine(businessdays[-1], self.business_hours[1])
return businessdays
def businesstimedelta(self, d1, d2):
"""
Returns a datetime.timedelta with business time between d1 and d2
"""
businessdays = self._build_spanning_datetimes(d1, d2)
time = datetime.timedelta()
if len(businessdays) == 0:
# HACK: manually handle the case when d1 is after business hours while d2 is during
if self.isduringbusinesshours(d2):
time += d2 - datetime.datetime.combine(d2, self.business_hours[0])
elif d2.time() > self.business_hours[1] and d1.time() < self.business_hours[0]:
time += self.open_hours
else:
prev = None
current = None
count = 0
for d in businessdays:
if current is None:
current = d
current = datetime.datetime.combine(d, current.time())
if prev is not None:
if prev.date() != current.date():
# time += datetime.timedelta(days=1)
time += self.open_hours
if count == len(businessdays) - 1:
if current > d:
# We went too far
# time -= datetime.timedelta(days=1)
time -= self.open_hours
time += self.open_hours - (current - d)
else:
time += d - current
count += 1
prev = current
return time