gaston.py

import openbabel as ob
import os.path,sys,re
from tempfile import mkstemp

def to_gasfile(input,type="sdf"):
    """ convert file to gaston format"""

    (n,gastonfile) = mkstemp()
    f = open(gastonfile, 'w')

    obc = ob.OBConversion()
    obc.SetInFormat(type)
    mol = ob.OBMol()
    next = obc.ReadFile(mol,input)
    molnum = 0

    while next:
        f.write("t # %d\n" % molnum)
        for i,atom in enumerate(ob.OBMolAtomIter(mol)): 
            f.write("v %d %d\n" % (i,atom.GetAtomicNum()))
        for i,bond in enumerate(ob.OBMolBondIter(mol)): 
            f.write("e %d %d %d\n" % (bond.GetBeginAtomIdx()-1,bond.GetEndAtomIdx()-1,bond.GetBondOrder()))

        mol = ob.OBMol()
        next = obc.Read(mol)
        molnum += 1
    return gastonfile

def from_gasfile(gastonfile,type="smi"):
    txt = open(gastonfile).read()
    p = re.compile('#.+?(?=(#|$))',re.S)
    mols = p.finditer(txt)

    for gas in mols:
        yield gas2ob(gas.group())
    os.unlink(gastonfile)


def gas2ob(gf):
    mol = ob.OBMol()

    for l in gf.split('\n'):
        if len(l) > 0 and l[0] == 'v':
            a = mol.NewAtom()
            atomic_num = int(l.split(' ')[2])
            a.SetAtomicNum(atomic_num)
        elif len(l) > 0 and l[0] == 'e':
            begin_atom_idx = int(l.split(' ')[1]) + 1
            end_atom_idx   = int(l.split(' ')[2]) + 1
            bond_order     = int(l.split(' ')[3])
            b = mol.AddBond(begin_atom_idx, end_atom_idx, bond_order)
        elif len(l) > 0 and l[0] == '#':
            title = l.split(' ')[1]
            mol.SetTitle(title)
    return mol

def do_gaston(gasfile,freq=7):
    (n,output) = mkstemp()
    os.system("/opt/local/bin/gaston %d %s %s" % (freq,gasfile,output))
    return output

def gaston(input):
    gasfile = to_gasfile(input)
    output = do_gaston(gasfile)
    os.unlink(gasfile)
    return from_gasfile(output)

if __name__ == '__main__':
    input = sys.argv[1]
    obc = ob.OBConversion()
    obc.SetOutFormat('smi')
    for m in gaston(input):
        sys.stdout.write(obc.WriteString(m))