Chem.ChemUtils.SDFToCSV

Package Chem :: Package ChemUtils :: Module SDFToCSV |
[hide private]
[frames] | no frames]

# 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