Support cmd.iterate expression as a decorated function (#449)

Author: pslacerda

modules/pymol/editing.py

@@ -66,6 +66,7 @@ def _iterate_prepare_args(expression, space, _self):
 if True:
 
     import math
+    from functools import partial
     from . import selector
     cmd = __import__("sys").modules["pymol.cmd"]
     from .cmd import _cmd,lock,unlock,is_string, \
@@ -1487,8 +1488,7 @@ def alter_list(object, expr_list, quiet=1, space=None, _self=cmd):
             return _cmd.alter_list(_self._COb, object, list(expr_list),
                                    int(quiet), dict(space))
 
-
-    def iterate(selection, expression, quiet=1, space=None, _self=cmd):
+    def iterate(selection, expression=None, quiet=1, space=None, _self=cmd):
 
         '''
 DESCRIPTION
@@ -1524,6 +1524,9 @@ def iterate(selection, expression, quiet=1, space=None, _self=cmd):
 
     iterate_state, alter, alter_state
         '''
+        if expression is None:
+            return partial(iterate, selection, quiet=quiet, space=space, _self=_self)
+
         expression, space = _iterate_prepare_args(expression, space, _self)
 
         # preprocess selection

testing/tests/api/editing.py

@@ -56,6 +56,16 @@ def testIterateCallback(self):
         self.assertEqual([a.name for a in atoms], ['N', 'CA', 'C', 'O'])
         self.assertEqual([a.protons for a in atoms], [7, 6, 6, 8])
 
+    @testing.requires_version("3.1")
+    def testIterateDecorator(self):
+        cmd.fragment('gly')
+        atoms = []
+        @cmd.iterate('not hydro')
+        def real_iterate(at):
+            atoms.append(at)
+        self.assertEqual([a.name for a in atoms], ['N', 'CA', 'C', 'O'])
+        self.assertEqual([a.protons for a in atoms], [7, 6, 6, 8])
+
     @testing.requires_version("2.6")
     def testIterate__explicit_degree(self):
         cmd.fragment("ala")