Completed
Push — master ( 6106cb...5ecb27 )
by Eddie
01:29
created

zipline.utils.get_early_closes()   B

Complexity

Conditions 4

Size

Total Lines 105

Duplication

Lines 0
Ratio 0 %
Metric Value
cc 4
dl 0
loc 105
rs 8.1935

How to fix   Long Method   

Long Method

Small methods make your code easier to understand, in particular if combined with a good name. Besides, if your method is small, finding a good name is usually much easier.

For example, if you find yourself adding comments to a method's body, this is usually a good sign to extract the commented part to a new method, and use the comment as a starting point when coming up with a good name for this new method.

Commonly applied refactorings include:

1
#
2
# Copyright 2013 Quantopian, Inc.
3
#
4
# Licensed under the Apache License, Version 2.0 (the "License");
5
# you may not use this file except in compliance with the License.
6
# You may obtain a copy of the License at
7
#
8
#     http://www.apache.org/licenses/LICENSE-2.0
9
#
10
# Unless required by applicable law or agreed to in writing, software
11
# distributed under the License is distributed on an "AS IS" BASIS,
12
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13
# See the License for the specific language governing permissions and
14
# limitations under the License.
15
import pandas as pd
16
import pytz
17
18
from datetime import datetime
19
from dateutil import rrule
20
from functools import partial
21
22
start = pd.Timestamp('1990-01-01', tz='UTC')
23
end_base = pd.Timestamp('today', tz='UTC')
24
# Give an aggressive buffer for logic that needs to use the next trading
25
# day or minute.
26
end = end_base + pd.Timedelta(days=365)
27
28
29
def canonicalize_datetime(dt):
30
    # Strip out any HHMMSS or timezone info in the user's datetime, so that
31
    # all the datetimes we return will be 00:00:00 UTC.
32
    return datetime(dt.year, dt.month, dt.day, tzinfo=pytz.utc)
33
34
35
def get_non_trading_days(start, end):
36
    non_trading_rules = []
37
38
    start = canonicalize_datetime(start)
39
    end = canonicalize_datetime(end)
40
41
    weekends = rrule.rrule(
42
        rrule.YEARLY,
43
        byweekday=(rrule.SA, rrule.SU),
44
        cache=True,
45
        dtstart=start,
46
        until=end
47
    )
48
    non_trading_rules.append(weekends)
49
50
    new_years = rrule.rrule(
51
        rrule.MONTHLY,
52
        byyearday=1,
53
        cache=True,
54
        dtstart=start,
55
        until=end
56
    )
57
    non_trading_rules.append(new_years)
58
59
    new_years_sunday = rrule.rrule(
60
        rrule.MONTHLY,
61
        byyearday=2,
62
        byweekday=rrule.MO,
63
        cache=True,
64
        dtstart=start,
65
        until=end
66
    )
67
    non_trading_rules.append(new_years_sunday)
68
69
    mlk_day = rrule.rrule(
70
        rrule.MONTHLY,
71
        bymonth=1,
72
        byweekday=(rrule.MO(+3)),
73
        cache=True,
74
        dtstart=datetime(1998, 1, 1, tzinfo=pytz.utc),
75
        until=end
76
    )
77
    non_trading_rules.append(mlk_day)
78
79
    presidents_day = rrule.rrule(
80
        rrule.MONTHLY,
81
        bymonth=2,
82
        byweekday=(rrule.MO(3)),
83
        cache=True,
84
        dtstart=start,
85
        until=end
86
    )
87
    non_trading_rules.append(presidents_day)
88
89
    good_friday = rrule.rrule(
90
        rrule.DAILY,
91
        byeaster=-2,
92
        cache=True,
93
        dtstart=start,
94
        until=end
95
    )
96
    non_trading_rules.append(good_friday)
97
98
    memorial_day = rrule.rrule(
99
        rrule.MONTHLY,
100
        bymonth=5,
101
        byweekday=(rrule.MO(-1)),
102
        cache=True,
103
        dtstart=start,
104
        until=end
105
    )
106
    non_trading_rules.append(memorial_day)
107
108
    july_4th = rrule.rrule(
109
        rrule.MONTHLY,
110
        bymonth=7,
111
        bymonthday=4,
112
        cache=True,
113
        dtstart=start,
114
        until=end
115
    )
116
    non_trading_rules.append(july_4th)
117
118
    july_4th_sunday = rrule.rrule(
119
        rrule.MONTHLY,
120
        bymonth=7,
121
        bymonthday=5,
122
        byweekday=rrule.MO,
123
        cache=True,
124
        dtstart=start,
125
        until=end
126
    )
127
    non_trading_rules.append(july_4th_sunday)
128
129
    july_4th_saturday = rrule.rrule(
130
        rrule.MONTHLY,
131
        bymonth=7,
132
        bymonthday=3,
133
        byweekday=rrule.FR,
134
        cache=True,
135
        dtstart=start,
136
        until=end
137
    )
138
    non_trading_rules.append(july_4th_saturday)
139
140
    labor_day = rrule.rrule(
141
        rrule.MONTHLY,
142
        bymonth=9,
143
        byweekday=(rrule.MO(1)),
144
        cache=True,
145
        dtstart=start,
146
        until=end
147
    )
148
    non_trading_rules.append(labor_day)
149
150
    thanksgiving = rrule.rrule(
151
        rrule.MONTHLY,
152
        bymonth=11,
153
        byweekday=(rrule.TH(4)),
154
        cache=True,
155
        dtstart=start,
156
        until=end
157
    )
158
    non_trading_rules.append(thanksgiving)
159
160
    christmas = rrule.rrule(
161
        rrule.MONTHLY,
162
        bymonth=12,
163
        bymonthday=25,
164
        cache=True,
165
        dtstart=start,
166
        until=end
167
    )
168
    non_trading_rules.append(christmas)
169
170
    christmas_sunday = rrule.rrule(
171
        rrule.MONTHLY,
172
        bymonth=12,
173
        bymonthday=26,
174
        byweekday=rrule.MO,
175
        cache=True,
176
        dtstart=start,
177
        until=end
178
    )
179
    non_trading_rules.append(christmas_sunday)
180
181
    # If Christmas is a Saturday then 24th, a Friday is observed.
182
    christmas_saturday = rrule.rrule(
183
        rrule.MONTHLY,
184
        bymonth=12,
185
        bymonthday=24,
186
        byweekday=rrule.FR,
187
        cache=True,
188
        dtstart=start,
189
        until=end
190
    )
191
    non_trading_rules.append(christmas_saturday)
192
193
    non_trading_ruleset = rrule.rruleset()
194
195
    for rule in non_trading_rules:
196
        non_trading_ruleset.rrule(rule)
197
198
    non_trading_days = non_trading_ruleset.between(start, end, inc=True)
199
200
    # Add September 11th closings
201
    # http://en.wikipedia.org/wiki/Aftermath_of_the_September_11_attacks
202
    # Due to the terrorist attacks, the stock market did not open on 9/11/2001
203
    # It did not open again until 9/17/2001.
204
    #
205
    #    September 2001
206
    # Su Mo Tu We Th Fr Sa
207
    #                    1
208
    #  2  3  4  5  6  7  8
209
    #  9 10 11 12 13 14 15
210
    # 16 17 18 19 20 21 22
211
    # 23 24 25 26 27 28 29
212
    # 30
213
214
    for day_num in range(11, 17):
215
        non_trading_days.append(
216
            datetime(2001, 9, day_num, tzinfo=pytz.utc))
217
218
    # Add closings due to Hurricane Sandy in 2012
219
    # http://en.wikipedia.org/wiki/Hurricane_sandy
220
    #
221
    # The stock exchange was closed due to Hurricane Sandy's
222
    # impact on New York.
223
    # It closed on 10/29 and 10/30, reopening on 10/31
224
    #     October 2012
225
    # Su Mo Tu We Th Fr Sa
226
    #     1  2  3  4  5  6
227
    #  7  8  9 10 11 12 13
228
    # 14 15 16 17 18 19 20
229
    # 21 22 23 24 25 26 27
230
    # 28 29 30 31
231
232
    for day_num in range(29, 31):
233
        non_trading_days.append(
234
            datetime(2012, 10, day_num, tzinfo=pytz.utc))
235
236
    # Misc closings from NYSE listing.
237
    # http://www.nyse.com/pdfs/closings.pdf
238
    #
239
    # National Days of Mourning
240
    # - President Richard Nixon
241
    non_trading_days.append(datetime(1994, 4, 27, tzinfo=pytz.utc))
242
    # - President Ronald W. Reagan - June 11, 2004
243
    non_trading_days.append(datetime(2004, 6, 11, tzinfo=pytz.utc))
244
    # - President Gerald R. Ford - Jan 2, 2007
245
    non_trading_days.append(datetime(2007, 1, 2, tzinfo=pytz.utc))
246
247
    non_trading_days.sort()
248
    return pd.DatetimeIndex(non_trading_days)
249
250
non_trading_days = get_non_trading_days(start, end)
251
trading_day = pd.tseries.offsets.CDay(holidays=non_trading_days)
252
253
254
def get_trading_days(start, end, trading_day=trading_day):
255
    return pd.date_range(start=start.date(),
256
                         end=end.date(),
257
                         freq=trading_day).tz_localize('UTC')
258
259
trading_days = get_trading_days(start, end)
260
261
262
def get_early_closes(start, end):
263
    # 1:00 PM close rules based on
264
    # http://quant.stackexchange.com/questions/4083/nyse-early-close-rules-july-4th-and-dec-25th # noqa
265
    # and verified against http://www.nyse.com/pdfs/closings.pdf
266
267
    # These rules are valid starting in 1993
268
269
    start = canonicalize_datetime(start)
270
    end = canonicalize_datetime(end)
271
272
    start = max(start, datetime(1993, 1, 1, tzinfo=pytz.utc))
273
    end = max(end, datetime(1993, 1, 1, tzinfo=pytz.utc))
274
275
    # Not included here are early closes prior to 1993
276
    # or unplanned early closes
277
278
    early_close_rules = []
279
280
    day_after_thanksgiving = rrule.rrule(
281
        rrule.MONTHLY,
282
        bymonth=11,
283
        # 4th Friday isn't correct if month starts on Friday, so restrict to
284
        # day range:
285
        byweekday=(rrule.FR),
286
        bymonthday=range(23, 30),
287
        cache=True,
288
        dtstart=start,
289
        until=end
290
    )
291
    early_close_rules.append(day_after_thanksgiving)
292
293
    christmas_eve = rrule.rrule(
294
        rrule.MONTHLY,
295
        bymonth=12,
296
        bymonthday=24,
297
        byweekday=(rrule.MO, rrule.TU, rrule.WE, rrule.TH),
298
        cache=True,
299
        dtstart=start,
300
        until=end
301
    )
302
    early_close_rules.append(christmas_eve)
303
304
    friday_after_christmas = rrule.rrule(
305
        rrule.MONTHLY,
306
        bymonth=12,
307
        bymonthday=26,
308
        byweekday=rrule.FR,
309
        cache=True,
310
        dtstart=start,
311
        # valid 1993-2007
312
        until=min(end, datetime(2007, 12, 31, tzinfo=pytz.utc))
313
    )
314
    early_close_rules.append(friday_after_christmas)
315
316
    day_before_independence_day = rrule.rrule(
317
        rrule.MONTHLY,
318
        bymonth=7,
319
        bymonthday=3,
320
        byweekday=(rrule.MO, rrule.TU, rrule.TH),
321
        cache=True,
322
        dtstart=start,
323
        until=end
324
    )
325
    early_close_rules.append(day_before_independence_day)
326
327
    day_after_independence_day = rrule.rrule(
328
        rrule.MONTHLY,
329
        bymonth=7,
330
        bymonthday=5,
331
        byweekday=rrule.FR,
332
        cache=True,
333
        dtstart=start,
334
        # starting in 2013: wednesday before independence day
335
        until=min(end, datetime(2012, 12, 31, tzinfo=pytz.utc))
336
    )
337
    early_close_rules.append(day_after_independence_day)
338
339
    wednesday_before_independence_day = rrule.rrule(
340
        rrule.MONTHLY,
341
        bymonth=7,
342
        bymonthday=3,
343
        byweekday=rrule.WE,
344
        cache=True,
345
        # starting in 2013
346
        dtstart=max(start, datetime(2013, 1, 1, tzinfo=pytz.utc)),
347
        until=max(end, datetime(2013, 1, 1, tzinfo=pytz.utc))
348
    )
349
    early_close_rules.append(wednesday_before_independence_day)
350
351
    early_close_ruleset = rrule.rruleset()
352
353
    for rule in early_close_rules:
354
        early_close_ruleset.rrule(rule)
355
    early_closes = early_close_ruleset.between(start, end, inc=True)
356
357
    # Misc early closings from NYSE listing.
358
    # http://www.nyse.com/pdfs/closings.pdf
359
    #
360
    # New Year's Eve
361
    nye_1999 = datetime(1999, 12, 31, tzinfo=pytz.utc)
362
    if start <= nye_1999 and nye_1999 <= end:
363
        early_closes.append(nye_1999)
364
365
    early_closes.sort()
366
    return pd.DatetimeIndex(early_closes)
367
368
early_closes = get_early_closes(start, end)
369
370
371
def get_open_and_close(day, early_closes):
372
    market_open = pd.Timestamp(
373
        datetime(
374
            year=day.year,
375
            month=day.month,
376
            day=day.day,
377
            hour=9,
378
            minute=31),
379
        tz='US/Eastern').tz_convert('UTC')
380
    # 1 PM if early close, 4 PM otherwise
381
    close_hour = 13 if day in early_closes else 16
382
    market_close = pd.Timestamp(
383
        datetime(
384
            year=day.year,
385
            month=day.month,
386
            day=day.day,
387
            hour=close_hour),
388
        tz='US/Eastern').tz_convert('UTC')
389
390
    return market_open, market_close
391
392
393
def get_open_and_closes(trading_days, early_closes, get_open_and_close):
394
    open_and_closes = pd.DataFrame(index=trading_days,
395
                                   columns=('market_open', 'market_close'))
396
397
    get_o_and_c = partial(get_open_and_close, early_closes=early_closes)
398
399
    open_and_closes['market_open'], open_and_closes['market_close'] = \
400
        zip(*open_and_closes.index.map(get_o_and_c))
401
402
    return open_and_closes
403
404
open_and_closes = get_open_and_closes(trading_days, early_closes,
405
                                      get_open_and_close)
406