Source Code for Module QSTK.csvconverter.yahoo_csv_to_pkl

  1  ''' 
  2  Created on Feb 25, 2011 
  3   
  4  @note: This assumes that all the CSV files will have exactly 7 columns. It ignores the first row in the csv files because it it the heading. 
  5          All the data is converted from csv to pkl 
  6  @author: Shreyas Joshi 
  7  @contact: shreyasj@gatech.edu 
  8  @summary: This is used to convert CSV files from norgate to pkl files- which are used when the simulator runs. 
  9  ''' 
 10   
 11  import numpy as np 
 12  import pickle as pkl 
 13  import qstkutil.utils as utils 
 14  import os 
 15  import dircache 
 16  import time 
 17  import sys 
 18   
 19   
 20 -def main (): 
 21       
 22      print "Starting..."+ str(time.strftime("%H:%M:%S")) 
 23       
 24      try: 
 25          rootdir = os.environ['QSDATA'] 
 26      except KeyError: 
 27          #rootdir = "/hzr71/research/QSData" 
 28          print "Please be sure to set the value for QSDATA in config.sh or local.sh\n"     
 29       
 30      fileExtensionToRemove = ".csv" 
 31       
 32      listOfInputPaths= list() 
 33   
 34       
 35  #For Gekko 
 36      #listOfInputPaths.append("/hzr71/research/QSData/Processed/Norgate/raw/Delisted Securities/US Recent/") 
 37      listOfInputPaths.append (rootdir + "/Raw/Yahoo/US/AMEX/") 
 38      listOfInputPaths.append (rootdir + "/Raw/Yahoo/US/NASDAQ/") 
 39      listOfInputPaths.append (rootdir + "/Raw/Yahoo/US/NYSE/") 
 40       
 41       
 42       
 43      listOfOutputPaths= list() 
 44  #    listOfOutputPaths.append("C:\\test\\temp\\pkl1\\") 
 45  #    listOfOutputPaths.append("C:\\test\\temp\\pkl2\\")     
 46       
 47      #listOfOutputPaths.append(rootdir + "/Norgate/Delisted Securities/US Recent/") 
 48      listOfOutputPaths.append(rootdir + "/Processed/Yahoo/US/AMEX/") 
 49      listOfOutputPaths.append(rootdir + "/Processed/Yahoo/US/NASDAQ/") 
 50      listOfOutputPaths.append(rootdir + "/Processed/Yahoo/US/NYSE/")     
 51       
 52      #If the output paths don't exist, then create them... 
 53      for path in listOfOutputPaths: 
 54          if not (os.access(path, os.F_OK)): 
 55              #Path does not exist, so create it 
 56              os.makedirs(path) #Makes paths recursively 
 57      #done making all output paths! 
 58       
 59      #In case there are already some files there- remove them. This will remove all the pkl fils from the previous run 
 60      utils.clean_paths (listOfOutputPaths) 
 61       
 62       
 63      if (len(listOfInputPaths)!= len(listOfOutputPaths)): 
 64          print "No. of input paths not equal to the number of output paths.. quitting" 
 65          sys.exit("FAILURE") 
 66          #if ends 
 67       
 68      path_ctr = -1; 
 69      use_cols = range (1, 7 + 1) # will now use cols 1 to 7 
 70      for path in listOfInputPaths: 
 71          path_ctr =  path_ctr + 1; 
 72          stocks_at_this_path = dircache.listdir(str(path)) 
 73          #Next, throw away everything that is not a .csv And these are our stocks! Example: this should throw away the '$' folder in the NYSE folder 
 74          filtered_names= filter (lambda x:(str(x).find(str(fileExtensionToRemove)) > -1), stocks_at_this_path) 
 75          #Now, we remove the .csv to get the name of the stock 
 76          filtered_names = map(lambda x:(x.partition(str(fileExtensionToRemove))[0]),filtered_names) 
 77          stock_ctr = -1 
 78          for stock in filtered_names: 
 79              stock_ctr = stock_ctr + 1 
 80              print "Reading file: " + str (path + stock) 
 81              #read in the stock date from the CSV file 
 82              stock_data= np.loadtxt (path + stock+".csv", np.float, None, ",", None, 1, use_cols) 
 83               
 84              stock_data_shape = stock_data.shape 
 85              #print "stock_data_shape is: " + str(stock_data_shape) 
 86           
 87               
 88  #            for i in range (0, stock_data_shape[0]): 
 89  #                print stock_data [i,: ] 
 90               
 91  #            print "Reading: " + str(stock) 
 92              f = open (listOfOutputPaths[path_ctr] + filtered_names[stock_ctr] + ".pkl", "wb" ) 
 93              pkl.dump (stock_data, f, -1) 
 94              f.close() 
 95          #for stock in stocks_at_this_path ends 
 96      #for path in listOfInputPaths ends 
 97      print "Finished..."+ str(time.strftime("%H:%M:%S")) 
 98       
 99      #main ends 
100       
101       
102  if __name__ == '__main__': 
103      main() 
104