Source Code for Module helpers.TransposeCsv

  1  #! /usr/bin/env python 
  2  #------------------------------------------------------------------------------ 
  3  #$Id: TransposeCsv.py 7080 2010-06-10 15:34:33Z EckhardSutorius $ 
  4  """ 
  5     Transposes given CSV file. 
  6   
  7     @author: E. Sutorius 
  8     @org:    WFAU, IfA, University of Edinburgh 
  9   
 10     @newfield contributors: Contributors, Contributors (Alphabetical Order) 
 11     @contributors: R.S. Collins 
 12   
 13     @todo: Make use of utils.extractColumns() when then gets CSV support. 
 14  """ 
 15  #------------------------------------------------------------------------------ 
 16  import os 
 17   
 18  from   wsatools.CLI                 import CLI 
 19  import wsatools.CSV                     as csv 
 20  from   wsatools.DbConnect.DbSession import DbSession 
 21  from   wsatools.Logger              import Logger, ForLoopMonitor 
 22  #------------------------------------------------------------------------------ 
 23  # Entry point for script. 
 24   
 25  # Allow module to be imported as well as executed from the command line 
 26  if __name__ == '__main__': 
 27      outDef = 'ingest.csv.tpd' 
 28      CLI.progArgs.append(CLI.Argument('input', 'ingest.csv')) 
 29      CLI.progOpts.remove('test') 
 30      CLI.progOpts.remove('user') 
 31      CLI.progOpts += [ 
 32        CLI.Option('b', 'begin', 
 33          'Initial row to read', 
 34          'NUMBER'), 
 35        CLI.Option('e', 'end', 
 36          'Final row to read', 
 37          'NUMBER'), 
 38        CLI.Option('a', 'attr', 
 39          'Attribute list', 
 40          'STR'), 
 41        CLI.Option('o', 'output', 
 42          'Alternative output file name', 
 43          'FILE', outDef), 
 44        CLI.Option('t', 'table', 
 45          'DB table name', 
 46          'NAME')] 
 47   
 48      cli = CLI("TransposeCSV", "$Revision: 7080 $", __doc__) 
 49      Logger.isVerbose = False 
 50      Logger.addMessage(cli.getProgDetails()) 
 51   
 52      if not os.path.exists(cli.getArg('input')): 
 53          Logger.addMessage( 
 54            "<ERROR> File does not exist: " + cli.getArg('input')) 
 55          raise SystemExit 
 56   
 57      fileName = os.path.basename(cli.getArg('input')) 
 58      db_table = (cli.getOpt('table') or 
 59        os.path.splitext(fileName)[0].partition('.')[0].rpartition('_')[2]) 
 60      if not db_table: 
 61          Logger.addMessage("No table specified!") 
 62          raise SystemExit 
 63   
 64      # Query database for table information 
 65      colNames = DbSession(cli.getArg('database')).queryColumnNames(db_table) 
 66      inputLines = csv.File(cli.getArg('input')).readlines() 
 67      if inputLines[0] == tuple(colNames): 
 68          inputLines.pop(0) 
 69   
 70      attrList = (cli.getOpt('attr').split(',') if cli.getOpt('attr') 
 71                  else colNames) 
 72   
 73      transDict = dict((col, [col]) for col in attrList) 
 74   
 75      if cli.getOpt('begin') and cli.getOpt('end'): 
 76          inputLines = \ 
 77            inputLines[int(cli.getOpt('begin')):int(cli.getOpt('end'))] 
 78      elif cli.getOpt('begin'): 
 79          inputLines = inputLines[int(cli.getOpt('begin')):] 
 80      elif cli.getOpt('end'): 
 81          inputLines = inputLines[:int(cli.getOpt('end'))] 
 82   
 83      progress = ForLoopMonitor(inputLines) 
 84      for line in inputLines: 
 85          for num, entry in enumerate(line): 
 86              if colNames[num] in attrList: 
 87                  transDict[colNames[num]].append(entry) 
 88          progress.testForOutput() 
 89   
 90      outPath = (cli.getOpt('output') if cli.getOpt('output') is not outDef else 
 91                 fileName + '.tpd') 
 92   
 93      csv.File(outPath, 'w').writelines(transDict[entry] for entry in attrList) 
 94   
 95      Logger.addMessage("Output written to %s" % outPath) 
 96   
 97  #------------------------------------------------------------------------------ 
 98  # Change log: 
 99  # 
100  # 29-Nov-2007, ETWS: Original version. 
101