Source Code for Module Chem.ChemUtils.SDFToCSV

  1  # $Id: SDFToCSV.py 2 2006-05-06 22:54:39Z glandrum $ 
  2  # 
  3  # Copyright (C) 2004-2006 Rational Discovery LLC 
  4  # 
  5  #   @@ All Rights Reserved  @@ 
  6  # 
  7  import RDConfig 
  8  import Chem 
  9  import sys,csv 
 10   
 11 -def Convert(suppl,outFile,keyCol='',stopAfter=-1,includeChirality=0,smilesFrom=''): 
 12    w = csv.writer(outFile) 
 13    mol = suppl[0] 
 14    propNames = list(mol.GetPropNames()) 
 15    if keyCol and keyCol in propNames: 
 16      propNames.remove(keyCol) 
 17   
 18    outL = [] 
 19    if keyCol: 
 20      outL.append(keyCol) 
 21    outL.append('SMILES') 
 22    outL.extend(propNames) 
 23    w.writerow(outL) 
 24    nDone = 0 
 25    for mol in suppl: 
 26      if not smilesFrom or not mol.HasProp(smilesFrom): 
 27        smi = Chem.MolToSmiles(mol,includeChirality) 
 28      else: 
 29        smi = mol.GetProp(smilesFrom) 
 30        tMol = Chem.MolFromSmiles(smi) 
 31        smi = Chem.MolToSmiles(tMol,includeChirality) 
 32      outL = [] 
 33      if keyCol: 
 34        outL.append(str(mol.GetProp(keyCol))) 
 35      outL.append(smi) 
 36      for prop in propNames: 
 37        if mol.HasProp(prop): 
 38          outL.append(str(mol.GetProp(prop))) 
 39        else: 
 40          outL.append('') 
 41      w.writerow(outL) 
 42      nDone += 1 
 43      if nDone == stopAfter: 
 44        break 
 45    return 
 46   
 47   
 48  #------------------- 
 49  #  Testing: 
 50  import unittest 
 51 -class TestCase(unittest.TestCase): 
 52 -  def setUp(self): 
 53      pass 
 54 -  def tearDown(self): 
 55      pass 
 56 -  def test1(self): 
 57      import os 
 58      from cStringIO import StringIO 
 59      fName = os.path.join(RDConfig.RDDataDir,'NCI','first_200.props.sdf') 
 60      suppl = Chem.SDMolSupplier(fName) 
 61      io = StringIO() 
 62      try: 
 63        Convert(suppl,io) 
 64      except: 
 65        import traceback 
 66        traceback.print_exc() 
 67        self.fail('conversion failed') 
 68      txt = io.getvalue() 
 69      lines = txt.split('\n') 
 70      if not lines[-1]: 
 71        del lines[-1] 
 72      self.failUnless(len(lines)==201,'bad num lines: %d'%len(lines)) 
 73      line0 = lines[0].split(',') 
 74      self.failUnless(len(line0)==16,'%d'%len(line0)) 
 75      self.failUnless(line0[0]=='SMILES') 
 76 -  def test2(self): 
 77      import os 
 78      from cStringIO import StringIO 
 79      fName = os.path.join(RDConfig.RDDataDir,'NCI','first_200.props.sdf') 
 80      suppl = Chem.SDMolSupplier(fName) 
 81      io = StringIO() 
 82      try: 
 83        Convert(suppl,io,keyCol='AMW',stopAfter=5) 
 84      except: 
 85        import traceback 
 86        traceback.print_exc() 
 87        self.fail('conversion failed') 
 88      txt = io.getvalue() 
 89      lines = txt.split('\n') 
 90      if not lines[-1]: 
 91        del lines[-1] 
 92      self.failUnless(len(lines)==6,'bad num lines: %d'%len(lines)) 
 93      line0 = lines[0].split(',') 
 94      self.failUnless(len(line0)==16,'%d'%len(line0)) 
 95      self.failUnless(line0[0]=='AMW') 
 96      self.failUnless(line0[1]=='SMILES') 
 97       
 98       
 99       
100     
101   
102   
103   
104  #------------------- 
105  #  CLI STuff: 
106 -def Usage(): 
107    message = """ 
108    Usage: SDFToCSV [-k keyCol] inFile.sdf [outFile.csv] 
109   
110    """ 
111    sys.stderr.write(message) 
112    sys.exit(-1) 
113   
114   
115   
116  if __name__=='__main__': 
117    import getopt 
118   
119    try: 
120      args,extras = getopt.getopt(sys.argv[1:],'hk:', 
121                                  ['test', 
122                                   'chiral', 
123                                   'smilesCol=', 
124                                   ]) 
125    except: 
126      import traceback 
127      traceback.print_exc() 
128      Usage() 
129   
130    keyCol = '' 
131    testIt = 0 
132    useChirality=0 
133    smilesCol='' 
134    for arg,val in args: 
135      if arg=='-k': 
136        keyCol = val 
137      elif arg=='--chiral': 
138        useChirality=1 
139      elif arg=='--smilesCol': 
140        smilesCol=val 
141      elif arg=='--test': 
142        testIt=1 
143      elif arg=='-h': 
144        Usage() 
145   
146    if not testIt and len(extras)<1: 
147      Usage() 
148         
149   
150    if not testIt: 
151      inFilename = extras[0] 
152      if len(extras)>1: 
153        outFilename = extras[1] 
154        outF = open(outFilename,'w+') 
155      else: 
156        outF = sys.stdout 
157   
158      suppl = Chem.SDMolSupplier(inFilename) 
159      Convert(suppl,outF,keyCol=keyCol,includeChirality=useChirality,smilesFrom=smilesCol) 
160    else: 
161      sys.argv = [sys.argv[0]] 
162      unittest.main() 
163