Fix merges.

Author: pslacerda1

modules/pymol/cmd.py

@@ -202,7 +202,7 @@ def as_pathstr(path):
 
         # for extending the language
 
-        from .commanding import extend, extendaa, alias
+        from .commanding import declare_command, extend, extendaa, alias
 
         # for documentation etc
 

modules/pymol/commanding.py

@@ -18,11 +18,23 @@
     if True:
         import _thread as thread
         import urllib.request as urllib2
-        from io import FileIO as file
+        from io import FileIO as file, BytesIO
+
+    import inspect
+    import glob
+    import shlex
+    import tokenize
+    from enum import Enum
+    from functools import wraps
+    from pathlib import Path
+    from textwrap import dedent
+    from typing import Tuple, Iterable, get_args,  Optional, Union, Any, NewType, List, get_origin
+
 
     import re
     import os
     import time
+    import builtins 
     import threading
     import traceback
     from . import colorprinting
@@ -529,6 +541,168 @@ def delete(name, *, _self=cmd):
         if _self._raising(r,_self): raise pymol.CmdException
         return r
 
+    # Selection = NewType('Selection', str)
+    
+    def _into_types(type, value):
+        if repr(type) == 'typing.Any':
+            return value
+        elif type is bool:
+            if isinstance(value, bool):
+                return value
+            if value.lower() in ["yes", "1", "true", "on", "y"]:
+                return True
+            elif value.lower() in ["no", "0", "false", "off", "n"]:
+                return False
+            else:
+                raise pymol.CmdException("Invalid boolean value: %s" % value)
+        
+        elif isinstance(type, builtins.type):
+            return type(value)
+
+        if origin := get_origin(type):
+            if not repr(origin).startswith('typing.') and issubclass(origin, tuple):
+                args = get_args(type)
+                new_values = []
+                for i, new_value in enumerate(shlex.split(value)):
+                    new_values.append(_into_types(args[i], new_value))
+                return tuple(new_values)
+            
+            elif origin == Union:
+                args = get_args(type)
+                found = False
+                for i, arg in enumerate(args):
+                    try:
+                        found = True
+                        return _into_types(arg, value)
+                    except:
+                        found = False
+                if not found:
+                    raise pymol.CmdException(f"Union was not able to cast %s" % value)
+                    
+            elif issubclass(list, origin):
+                args = get_args(type)
+                if len(args) > 0:
+                    f = args[0]
+                else:
+                    f = lambda x: x
+                return [f(i) for i in shlex.split(value)]
+        
+        
+        # elif value is None:
+        #     origin = get_origin(type)
+        #     if origin is None:
+        #         return None
+        #     else:
+        #         return _into_types(origin)
+        #         for arg in get_args(origin):
+        #         return _into_types(get_args(origin), value)
+        
+        elif isinstance(type, str):
+            return str(value)
+
+        raise pymol.CmdException(f"Unsupported argument type {type}")
+                    
+    def parse_documentation(func):
+        source = inspect.getsource(func)
+        tokens = tokenize.tokenize(BytesIO(source.encode('utf-8')).readline)
+        tokens = list(tokens)
+        comments = []
+        params = {}
+        i = -1
+        started = False
+        while True:
+            i += 1
+            if tokens[i].string == "def":
+                while tokens[i].string == "(":
+                    i += 1
+                started = True
+                continue
+            if not started:
+                continue
+            if tokens[i].string == "->":
+                break
+            if tokens[i].type == tokenize.NEWLINE:
+                break
+            if tokens[i].string == ")":
+                break
+            if tokens[i].type == tokenize.COMMENT:
+                comments.append(tokens[i].string)
+                continue
+            if tokens[i].type == tokenize.NAME and tokens[i+1].string == ":":
+                name = tokens[i].string
+                name_line = tokens[i].line
+                i += 1
+                while not (tokens[i].type == tokenize.NAME and tokens[i+1].string == ":"):
+                    if tokens[i].type == tokenize.COMMENT and tokens[i].line == name_line:
+                        comments.append(tokens[i].string)
+                        break
+                    elif tokens[i].type == tokenize.NEWLINE:
+                        break
+                    i += 1
+                else:
+                    i -= 3
+                docs = ' '.join(c[1:].strip() for c in comments)
+                params[name] = docs
+                comments = []
+        return params
+
+
+    def declare_command(name, function=None, _self=cmd):
+
+        if function is None:
+            name, function = name.__name__, name
+
+        # docstring text, if present, should be dedented
+        if function.__doc__ is not None:
+            function.__doc__ = dedent(function.__doc__)
+
+        # Analysing arguments
+        spec = inspect.getfullargspec(function)
+        kwargs_ = {}
+        args_ = spec.args[:]
+        defaults = list(spec.defaults or [])
+
+        args2_ = args_[:]
+        while args_ and defaults:
+            kwargs_[args_.pop(-1)] = defaults.pop(-1)
+
+        funcs = {}
+        for idx, (var, func) in enumerate(spec.annotations.items()):
+            funcs[var] = func
+
+        # Inner function that will be callable every time the command is executed
+        @wraps(function)
+        def inner(*args, **kwargs):
+            frame = traceback.format_stack()[-2]
+            caller = frame.split("\"", maxsplit=2)[1]
+            # It was called from command line or pml script, so parse arguments
+            if caller.endswith("pymol/parser.py"):
+                kwargs = {**kwargs, **dict(zip(args2_, args))}
+                kwargs.pop("_self", None)
+                new_kwargs = {}
+                for var, type in funcs.items():
+                    if var in kwargs:
+                        value = kwargs[var]
+                        new_kwargs[var] = _into_types(type, value)
+                final_kwargs = {}
+                for k, v in kwargs_.items():
+                    final_kwargs[k] = v
+                for k, v in new_kwargs.items():
+                    if k not in final_kwargs:
+                        final_kwargs[k] = v
+                return function(**final_kwargs)
+
+            # It was called from Python, so pass the arguments as is
+            else:
+                return function(*args, **kwargs)
+        inner.__arg_docs = parse_documentation(function)
+
+        _self.keyword[name] = [inner, 0,0,',',parsing.STRICT]
+        _self.kwhash.append(name)
+        _self.help_sc.append(name)
+        return inner
+
+
     def extend(name, function=None, _self=cmd):
 
         '''

testing/tests/api/commanding.py

@@ -1,9 +1,13 @@
 from __future__ import print_function
 
 import sys
-import pymol
+
 import __main__
 from pymol import cmd, testing, stored
+from typing import List
+from typing import Optional, Any, Tuple, Union
+from pathlib import Path
+
 
 class TestCommanding(testing.PyMOLTestCase):
 
@@ -171,3 +175,137 @@ def testRun(self, namespace, mod, rw):
             self.assertTrue(stored.tmp)
             if mod:
                 self.assertEqual(rw, hasattr(sys.modules[mod], varname))
+
+
+def test_declare_command_casting():
+    from pathlib import Path
+
+
+    @cmd.declare_command
+    def func(a: int, b: Path):
+        assert isinstance(a, int) and a == 1
+        assert isinstance(b, Path) and "/tmp" == str(b)
+    cmd.do('func 1, /tmp')
+
+def test_declare_command_optional(capsys):
+    @cmd.declare_command
+    def func(a: Optional[int] = None):
+        assert a is None
+    cmd.do("func")
+    out, err = capsys.readouterr()
+    assert out+err == ''
+
+    @cmd.declare_command
+    def func(a: Optional[int] = None):
+        assert a is 10
+    cmd.do("func 10")
+    out, err = capsys.readouterr()
+    assert out+err == ''
+
+def test_declare_command_docstring():
+    @cmd.declare_command
+    def func():
+        """docstring"""
+    assert func.__doc__ == "docstring"
+
+
+def test_declare_command_bool(capsys):
+    @cmd.declare_command
+    def func(a: bool, b: bool):
+        assert a
+        assert not b
+
+    cmd.do("func yes, 0")
+    out, err = capsys.readouterr()
+    assert out == '' and err == ''
+
+
+def test_declare_command_generic(capsys):
+    @cmd.declare_command
+    def func(
+        nullable_point: Tuple[float, float, float],
+        my_var: Union[int, float] = 10,
+        my_foo: Union[int, float] = 10.0,
+        extended_calculation: bool = True,
+        old_style: Any = "Old behavior"
+    ):
+        assert nullable_point == (1., 2., 3.)
+        assert extended_calculation
+        assert isinstance(my_var, int)
+        assert isinstance(my_foo, float)
+        assert old_style == "Old behavior"
+
+    cmd.do("func nullable_point=1 2 3, my_foo=11.0")
+    out, err = capsys.readouterr()
+    assert out + err == ''
+
+def test_declare_command_path(capsys):
+    @cmd.declare_command
+    def func(dirname: Path = Path('.')):
+        assert dirname.exists()
+    cmd.do('func ..')
+    cmd.do('func')
+    out, err = capsys.readouterr()
+    assert out + err == ''
+
+def test_declare_command_any(capsys):
+    @cmd.declare_command
+    def func(old_style: Any):
+        assert old_style != "RuntimeError"
+    cmd.do("func RuntimeError")
+    out, err = capsys.readouterr()
+    assert 'AssertionError' in out+err
+
+def test_declare_command_list(capsys):
+    @cmd.declare_command
+    def func(a: List):
+        assert a[1] == "2"
+    cmd.do("func 1 2 3")
+    out, err = capsys.readouterr()
+    assert out + err == ''
+
+    @cmd.declare_command
+    def func(a: List[int]):
+        assert a[1] == 2
+    cmd.do("func 1 2 3")
+    out, err = capsys.readouterr()
+    assert out + err == ''
+
+def test_declare_command_tuple(capsys):
+    @cmd.declare_command
+    def func(a: Tuple[str, int]):
+        assert a == ("fooo", 42)
+    cmd.do("func fooo 42")
+    out, err = capsys.readouterr()
+    assert out + err == ''
+
+def test_declare_command_arg_docs():
+    @cmd.declare_command
+    def func(
+        # multiline
+        # documentation works
+        foo: int, # inline
+        a: str,
+        # bar are strings
+        bar: Tuple[str, int], # continued...
+        b: Any = 10, # The new old age
+        # aaaa
+        c: Any = 'a' # b
+    ):
+        "main description"
+        pass
+
+    assert func.__arg_docs['foo'] == "multiline documentation works inline"
+    assert func.__arg_docs['a'] == ""
+    assert func.__arg_docs['bar'] == "bar are strings continued..."
+    assert func.__arg_docs['b'] == 'The new old age'
+    assert func.__arg_docs['c'] == 'aaaa b'
+    assert func.__annotations__['foo'] == int
+    assert func.__annotations__['bar'] == Tuple[str, int]
+
+def test_declare_command_default():
+    @cmd.declare_command
+    def func(a: str="sele"):
+        assert a == "a"
+    func("a")
+    cmd.do('func a')