Source Code for Module QSTK.csvconverter.compustat_csv_to_pkl

  1  ''' 
  2  Created on June 1, 2011 
  3   
  4  @author: John Cornwell 
  5  @contact: JohnWCornwellV@gmail.com 
  6  @summary: Used to extract Compustat data from a csv dump and convert into individual pickle files 
  7  ''' 
  8   
  9  import numpy as np 
 10  import datetime as dt 
 11  import pickle as pkl 
 12  import qstkutil.utils as utils 
 13  import qstkutil.DataAccess as da 
 14  import os 
 15  import dircache 
 16  import time 
 17  import csv 
 18  import sys 
 19   
 20   
 21 -def _dumpFiles( dData, lSets, lsOutPaths ): 
 22      ''' 
 23      @summary Helper function to store files on disk: files from dictionary dData, set of paths lSets 
 24      @param dData: Dictionary of indexes to data label strings.  
 25      @param lSets: List of symbol sets, each corresponds to a directory, e.g. NYSE, NASDAQ. 
 26      @param lsOutPaths: List of path strings, same indexes as lSets. 
 27      ''' 
 28      lKeys = dData.keys() 
 29        
 30      for key in lKeys: 
 31          for i, symSet in enumerate( lSets ): 
 32              if key in symSet: 
 33                  sFilename = lsOutPaths[i] + key + '.pkl' 
 34                  break 
 35           
 36          fOut = open( sFilename, 'wb' ) 
 37          pkl.dump( dData[key], fOut, -1) 
 38          fOut.close() 
 39       
 40       
 41 -def _analyze(): 
 42      ''' 
 43      @summary Helper function to analyzes the csv data to determine which columns are float values, and generate the set of valid labels. 
 44      ''' 
 45      try: 
 46          rootdir = os.environ['QSDATA'] 
 47      except KeyError: 
 48          print "Please be sure to set the value for QSDATA in config.sh or local.sh\n"     
 49           
 50      ''' Create lists of input and output paths ''' 
 51      fFile = ( rootdir + "/Raw/Compustat/Compustat.csv") 
 52         
 53      spamReader = csv.reader(open(fFile, 'rb'), delimiter=',') 
 54       
 55      ''' Take first row as the labels ''' 
 56      for row in spamReader: 
 57          lsLabels = row 
 58          break     
 59       
 60      used = [0] * len(lsLabels) 
 61      badSet = set() 
 62       
 63      ''' find non-float valued entities ''' 
 64      for i, row in enumerate(spamReader): 
 65          for j, elem in enumerate(row): 
 66              if used[j] == 1: 
 67                  continue 
 68               
 69              if elem == '': 
 70                  elem = 0.0 
 71               
 72              try:  
 73                  float(elem) 
 74              except: 
 75                  badSet.add( lsLabels[j] ) 
 76                  used[j] = 1 
 77           
 78          if( i % 10000 == 0 ): 
 79              print (i / 1378625.0)*100, '%'    
 80       
 81      print 'Bad (non float) labels:' 
 82      print badSet    
 83       
 84      for i,label in enumerate(lsLabels): 
 85          if label in badSet: 
 86              del lsLabels[i] 
 87       
 88      print '\n\nGood (float) labels:' 
 89      print lsLabels 
 90       
 91      return   
 92       
 93   
 94 -def convert (): 
 95      ''' 
 96      @summary: Converts a Compustat CSV file to pickle files of numpy arrays. 
 97      ''' 
 98       
 99      print "Starting..."+ str(time.strftime("%H:%M:%S")) 
100       
101      ''' Write every so often to save memory, 20k lines is usually < .5GB ''' 
102      lSaveMem = 20000 
103       
104      try: 
105          rootdir = os.environ['QSDATA'] 
106      except KeyError: 
107          print "Please be sure to set the value for QSDATA in config.sh or local.sh\n"     
108       
109      ''' Create lists of input and output paths ''' 
110      fFile = ( rootdir + "/Raw/Compustat/Compustat.csv") 
111   
112      listOfOutputPaths= [] 
113      listOfOutputPaths.append(rootdir + "/Processed/Compustat/US/AMEX/") 
114      listOfOutputPaths.append(rootdir + "/Processed/Compustat/US/NASDAQ/") 
115      listOfOutputPaths.append(rootdir + "/Processed/Compustat/US/NYSE/")     
116       
117      #If the output paths don't exist, then create them... 
118      for path in listOfOutputPaths: 
119          if not (os.access(path, os.F_OK)): 
120              #Path does not exist, so create it 
121              os.makedirs(path) #Makes paths recursively 
122      #done making all output paths! 
123       
124      #In case there are already some files there- remove them. This will remove all the pkl fils from the previous run 
125      utils.clean_paths (listOfOutputPaths) 
126         
127      spamReader = csv.reader(open(fFile, 'rb'), delimiter=',') 
128      ''' 1378625 rows typical '''     
129       
130      ''' Take first row as the labels ''' 
131      for row in spamReader: 
132          lsLabels = row 
133          break 
134      
135      ''' Generated from _Analyze() ''' 
136      lsBadLabels = set(['LOC', 'ADD4', 'ADD3', 'ADD2', 'ADD1', 'ACCTCHGQ', 'WEBURL', 'IDBFLAG', 'popsrc', 'DATACQTR', 'conm', 'COSTAT', 'FINALQ', 'fdateq', 'FAX', 'RP', 'PRIROW', 'dldte', 'indfmt', 'SPCSRC', 'BUSDESC', 'ipodate', 'PHONE', 'CURCDQ', 'pdateq', 'DATAFQTR', 'PRICAN', 'EIN', 'datadate', 'tic', 'ADDZIP', 'CONML', 'consol', 'datafmt', 'cusip', 'BSPRQ', 'OGMQ', 'COMPSTQ', 'COUNTY', 'STATE', 'CURNCDQ', 'CITY', 'rdq', 'apdedateq', 'STALTQ', 'INCORP']) 
137   
138      ''' get list of stocks in 3 US indexes ''' 
139      Access = da.DataAccess( 'Norgate' ) 
140      setNyse = set( Access.get_symbols_in_sublist("/US/NYSE") ) 
141      setNasdaq = set( Access.get_symbols_in_sublist("/US/NASDAQ") ) 
142      setAmex = set( Access.get_symbols_in_sublist("/US/AMEX") ) 
143       
144      ''' If stock appears in more than one index, remove to avoid ambiguity ''' 
145      print 'Ignoring duplicate stocks:', 
146      dup1 =  setNyse.intersection( setNasdaq.union(setAmex)) 
147      dup2 =  setNasdaq.intersection( setAmex ) 
148      print dup1.union(dup2) 
149   
150      setNyse   = setNyse - dup1.union(dup2) 
151      setAmex   = setAmex - dup1.union(dup2) 
152      setNasdaq = setNasdaq - dup1.union(dup2) 
153       
154      ''' Note the two lists below must be in the same order ''' 
155      lsOutPaths = [] 
156      lsOutPaths.append(rootdir + "/Processed/Compustat/US/AMEX/") 
157      lsOutPaths.append(rootdir + "/Processed/Compustat/US/NASDAQ/") 
158      lsOutPaths.append(rootdir + "/Processed/Compustat/US/NYSE/")   
159       
160      lSets = [setAmex, setNasdaq, setNyse]   
161       
162      #If the output paths don't exist, then create them... 
163      for path in lsOutPaths: 
164          if not (os.access(path, os.F_OK)): 
165              #Path does not exist, so create it 
166              os.makedirs(path) #Makes paths recursively 
167      #done making all output paths! 
168       
169      #In case there are already some files there- remove them. This will remove all the pkl fils from the previous run 
170      utils.clean_paths (lsOutPaths) 
171           
172      lDateCol = 0 
173      llUseCols = [] 
174   
175       
176      ''' We have first row (the labels), loop through saving good label indicies ''' 
177      for j, sElem in enumerate(lsLabels): 
178          if( sElem not in lsBadLabels ): 
179              llUseCols.append(j) 
180           
181          ''' Keep track of ticker column and date, specially handled later ''' 
182          if( sElem == 'datadate' ): 
183              lDateCol = j  
184          if( sElem == 'tic' ): 
185              lTicCol = j  
186           
187           
188      ''' Dict of ticker->numpy array mapping ''' 
189      dData = dict() 
190      print '' 
191       
192       
193      ''' Main loop, iterate over the rows in the csv file ''' 
194      for j, row in enumerate(spamReader): 
195          lsLabels = row 
196          sTic = row[lTicCol] 
197           
198          ''' Find out what index this belongs to ''' 
199          lIndex = -1 
200          for i, symSet in enumerate( lSets ): 
201              if sTic in symSet: 
202                  lIndex = i 
203                  break 
204          if lIndex == -1: 
205              continue 
206           
207          sFilename = lsOutPaths[lIndex] + sTic + '.pkl' 
208           
209          ''' If the file exists (temporary memory saving measure), read it in and delete file from disk ''' 
210          if( os.path.isfile(sFilename)): 
211              if dData.has_key(sTic): 
212                 print 'File should not be both on disk and in dict' 
213                 sys.exit("FAILURE") 
214               
215              fIn = open( sFilename, 'rb' ) 
216              dData[sTic] = pkl.load( fIn ) 
217              fIn.close() 
218              os.remove( sFilename ) 
219               
220          fDate = float( dt.datetime.strptime( row[lDateCol], "%m/%d/%Y").strftime("%Y%m%d") ) 
221           
222          ''' convert blanks to nans ''' 
223          for i in llUseCols: 
224              if row[i] == '': 
225                  row[i] = 'nan' 
226           
227          ''' Add row if data exists, if not, create new array ''' 
228          if dData.has_key(sTic):        
229              dData[sTic] = np.vstack( (dData[sTic], np.array([fDate] + [row[i] for i in llUseCols], dtype=np.float)) ) 
230          else: 
231              dData[sTic]= np.array( [fDate] + [row[i] for i in llUseCols], dtype=np.float ) 
232           
233          if( (j+1) % 1000 == 0): 
234              fDone = (j / 1378625.0) * 100 
235              print '\rApprox %.2lf%%'%((j / 1378625.0) * 100), 
236               
237          if( (j+1) % lSaveMem == 0): 
238              ''' Write all the pickle files we currently have ''' 
239               
240              print '\nWriting %i lines to pickle files to save memory\n'%(lSaveMem) 
241              _dumpFiles( dData, lSets, lsOutPaths) 
242              ''' Remember to delete! ''' 
243              del dData 
244              dData = dict() 
245               
246          # Done writing files 
247      # Done with main loop 
248   
249      print '' 
250      print 'Writing final pickle files\n'        
251      _dumpFiles( dData, lSets, lsOutPaths) 
252      del dData 
253       
254      print "Finished..."+ str(time.strftime("%H:%M:%S")) 
255       
256      return 
257       
258  if __name__ == '__main__': 
259   
260      #_analyze() 
261      convert() 
262