-
Notifications
You must be signed in to change notification settings - Fork 0
/
get_data.py
87 lines (68 loc) · 2.48 KB
/
get_data.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
import bs4 as bs
import datetime as dt
import os
import pandas as pd
import pandas_datareader.data as web
import pickle
import requests
import quandl
import csv
quandl.ApiConfig.api_key = 'Ri7BNHzdujUt3zYB7NQr'
def save_sp500_tickers():
resp = requests.get('http://en.wikipedia.org/wiki/List_of_S%26P_500_companies')
soup = bs.BeautifulSoup(resp.text, 'lxml')
table = soup.find('table', {'class': 'wikitable sortable'})
tickers = []
for row in table.findAll('tr')[1:]:
ticker = row.findAll('td')[0].text
tickers.append(ticker)
with open("sp500tickers.pickle","wb") as f:
pickle.dump(tickers,f)
return tickers
def get_data_from_quandl(reload_sp500=False):
tickers = []
if reload_sp500:
tickers = save_sp500_tickers()
else:
with open("sp500tickers.csv") as f:
tickers_csv = csv.reader(f)
for row in tickers_csv:
tickers.append(row[0])
#with open("sp500tickers.pickle","rb") as f:
# tickers = pickle.load(f)
print(tickers)
if not os.path.exists('stock_dfs'):
os.makedirs('stock_dfs')
start = dt.datetime(2000, 1, 1)
end = dt.datetime(2016, 12, 31)
for ticker in tickers:
# just in case your connection breaks, we'd like to save our progress!
if not os.path.exists('stock_dfs/{}.csv'.format(ticker)):
print(ticker)
df = quandl.get("WIKI/"+ticker, start_date= start, end_date= end)
df.to_csv('stock_dfs/{}.csv'.format(ticker))
else:
print('Already have {}'.format(ticker))
def compile_data():
tickers = []
with open("sp500tickers.csv") as f:
tickers_csv = csv.reader(f)
for row in tickers_csv:
tickers.append(row[0])
main_df = pd.DataFrame()
for count,ticker in enumerate(tickers):
df = pd.read_csv('stock_dfs/{}.csv'.format(ticker))
df.set_index('Date', inplace=True)
df.rename(columns={'Adj. Close':ticker}, inplace=True)
df.drop(['Open','High','Low','Close','Volume','Ex-Dividend','Split Ratio','Adj. Open','Adj. High','Adj. Low','Adj. Volume'],1,inplace=True)
if main_df.empty:
main_df = df
else:
main_df = main_df.join(df, how='outer')
if count % 10 == 0:
print(count)
print(main_df.head())
main_df.to_csv('sp500_joined_closes.csv')
#get_data_from_quandl(reload_sp500=False)
#save_sp500_tickers()
#compile_data()