Improvements on the @cmd.declare_command API

modules/pymol/commanding.py

@@ -21,16 +21,19 @@
     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 builtins
     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 List
+    from typing import Tuple, Iterable, get_args,  Optional, Union, Any, NewType, List, get_origin
+
 
     import re
     import os
@@ -600,45 +603,136 @@ def get_state_list(states_str):
             states_list = sorted(set(map(int, output)))
             return _cmd.delete_states(_self._COb, name, states_list)
 
-    class Selection(str):
-        pass
-
+    def _into_types(type, value):
+        """Convert a string value to an specific type."""
 
-    def _parse_bool(value: str):
-        if isinstance(value, str):
-            if value.lower() in ["yes", "1", "true", "on", "y"]:
-                return True
-            elif value.lower() in ["no", "0", "false", "off", "n"]:
-                return False
-            else:
-                raise Exception("Invalid boolean value: %s" % value)
-        elif isinstance(value, bool):
+        if repr(type) == 'typing.Any':
             return value
-        else:
-            raise Exception(f"Unsuported boolean flag {value}")
-
-    def _parse_list_str(value):
-        return shlex.split(value)
+
+        elif type is bool:
+            if isinstance(value, bool):
+                return value
+            elif isinstance(value, str):
+                if value.lower() in ["yes", "1", "true", "on", "y"]:
+                    return True
+                elif value.lower() in ["no", "0", "false", "off", "n"]:
+                    return False
+            elif isinstance(value, int):
+                return bool(value)
+            else:
+                raise pymol.CmdException(f"Invalid boolean value: {value}")
 
-    def _parse_list_int(value):
-        return list(map(int, shlex.split(value)))
+        elif isinstance(type, Enum):
+            if value in type:
+                return type(value)
+            else:
+                raise pymol.CmdException(f"Invalid value for enum {type.__name__}: {value}")
+
+        elif isinstance(type, builtins.type):
+            return type(value)
+
+        # Composite types for now
+        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("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)]
+
+        # TODO Optional/None case isn't working
+        # 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_args_docs(func):
+        """Extract the arguments documentation of a function.
+
+        They are given by the # comments preceding or at the same
+        line of each argument.
+        """
+        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 _parse_list_float(value):
-        return list(map(float, shlex.split(value)))
 
     def declare_command(name, function=None, _self=cmd):
+
         if function is None:
             name, function = name.__name__, name
 
-        # new style commands should have annotations
-        annotations = [a for a in function.__annotations__ if a != "return"]
-        if function.__code__.co_argcount != len(annotations):
-            raise Exception("Messy annotations")
-
         # docstring text, if present, should be dedented
         if function.__doc__ is not None:
-            function.__doc__ = dedent(function.__doc__).strip()
-
+            function.__doc__ = dedent(function.__doc__)
 
         # Analysing arguments
         spec = inspect.getfullargspec(function)
@@ -661,34 +755,24 @@ def inner(*args, **kwargs):
 
             # It was called from command line or pml script, so parse arguments
             if caller == _parser_filename:
-                kwargs = {**kwargs_, **kwargs, **dict(zip(args2_, args))}
+                kwargs = {**kwargs, **dict(zip(args2_, args))}
                 kwargs.pop("_self", None)
-                for arg in kwargs.copy():
-                    if funcs[arg] == bool:
-                        funcs[arg] = _parse_bool
-                    elif funcs[arg] == List[str]:
-                        funcs[arg] = _parse_list_str
-                    elif funcs[arg] == List[int]:
-                        funcs[arg] = _parse_list_int
-                    elif funcs[arg] == List[float]:
-                        funcs[arg] = _parse_list_float
-                    else:
-                        # Assume it's a literal supported type
-                        pass
-                    # Convert the argument to the correct type
-                    kwargs[arg] = funcs[arg](kwargs[arg])
-                return function(**kwargs)
+                new_kwargs = {}
+                for var, type in funcs.items():
+                    if var in kwargs:
+                        value = kwargs[var]
+                        new_kwargs[var] = _into_types(type, value)
+                return function(**new_kwargs)
 
             # It was called from Python, so pass the arguments as is
             else:
                 return function(*args, **kwargs)
+        inner.__arg_docs = parse_args_docs(function)
 
-        name = function.__name__
-        _self.keyword[name] = [inner, 0, 0, ",", parsing.STRICT]
-        _self.kwhash.append(name)
-        _self.help_sc.append(name)
+        _self.keyword[name] = [inner, 0,0,',',parsing.STRICT]
         return inner
 
+
     def extend(name, function=None, _self=cmd):
 
         '''