Pml2py

This script converts a pml script to a python script.

See pymol-users mailing list (Subject: Convert pml script to Pymol Python script, Fri, 8 Apr 2011).

import sys
from pymol import cmd

def pml2py(filename, out=sys.stdout):
    '''
DESCRIPTION

    Convert a pml script to python syntax.

USAGE

    pml2py infile [, outfile]

TODO

    comments, aliases, list arguments
    '''
    def quote(args):
        for arg in args:
            if '=' not in arg:
                yield repr(arg)
            else:
                a = arg.split('=', 1)
                yield a[0] + '=' + repr(a[1])
    if isinstance(out, basestring):
        out = open(out, 'w')
    print >> out, '''
# automatically converted from "%s"
import pymol
from pymol import *
''' % (filename)
    handle = iter(open(filename))
    for line in handle:
        while line.endswith('\\\n'):
            line = line[:-2] + handle.next()

        a = line.split(None, 1)
        try:
            name = a[0]
            if name.startswith('/'):
                line = line.lstrip()[1:]
                raise
            name = cmd.kwhash.shortcut.get(name, name)
            func = cmd.keyword[name][0]
            assert func.func_name != 'python_help'
        except:
            out.write(line)
            continue

        if func.func_name == '<lambda>':
            print >> out, 'cmd.do(%s)' % (repr(line.strip()))
            continue

        # FIXME hack for set_view list argument
        if name == 'set_view':
            print >> out, 'cmd.set_view(' + a[1] + ')'
            continue

        if len(a) > 1:
            args = [i.strip() for i in a[1].split(',')]
        else:
            args = []

        if name == 'python':
            for line in handle:
                if line.split() == ['python', 'end']:
                    break
                out.write(line)
            continue

        print >> out, '%s.%s(%s)' % (func.__module__, func.func_name, ', '.join(quote(args)))

cmd.extend('pml2py', pml2py)

# vi:expandtab:smarttab