diff --git a/modules/pymol/commanding.py b/modules/pymol/commanding.py index 3873c1baa..67a053198 100644 --- a/modules/pymol/commanding.py +++ b/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): ''' diff --git a/testing/tests/api/commanding.py b/testing/tests/api/commanding.py index b8bd4a541..11f0e6d95 100644 --- a/testing/tests/api/commanding.py +++ b/testing/tests/api/commanding.py @@ -1,15 +1,12 @@ -from __future__ import print_function - import sys -import pytest +from enum import Enum -import pymol import __main__ +from pytest import mark from pymol import cmd, testing, stored - from typing import List - - +from typing import Optional, Any, Tuple, Union, List +from pathlib import Path class TestCommanding(testing.PyMOLTestCase): @@ -185,26 +182,31 @@ def testRun(self, namespace, mod, rw): 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, str)) and "/tmp" == str(b) - func(1, "/tmp") + assert isinstance(b, Path) and "/tmp" == str(b) cmd.do('func 1, /tmp') -def test_declare_command_default(capsys): - from pymol.commanding import Selection +@mark.skip(reason="API not implemented yet") +def test_declare_command_optional(capsys): @cmd.declare_command - def func(a: Selection = "sele"): - assert a == "sele" - func() + def func(a: Optional[int] = None): + assert a is None cmd.do("func") out, err = capsys.readouterr() - assert out == '' + 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 @@ -212,72 +214,121 @@ def func(): """docstring""" assert func.__doc__ == "docstring" + +def test_declare_command_bool(capsys): @cmd.declare_command - def func(): - """ - docstring - Test: - --foo - """ - assert func.__doc__ == "docstring\nTest:\n --foo" + 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_type_return(capsys): - @cmd.declare_command - def func() -> int: - return 1 - assert func() == 1 +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 == '' + assert out + err == '' +def test_declare_command_path(capsys): @cmd.declare_command - def func(): - return 1 - assert func() == 1 + 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_list_str(capsys): +def test_declare_command_any(capsys): @cmd.declare_command - def func(a: List[str]): - print(a[-1]) + def func(old_style: Any): + assert old_style != "RuntimeError" + cmd.do("func RuntimeError") + out, err = capsys.readouterr() + assert 'AssertionError' in out+err - func(["a", "b", "c"]) - cmd.do('func a b c') +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 == 'c\nc\n' + assert out + err == '' -def test_declare_command_list_int(capsys): @cmd.declare_command def func(a: List[int]): - print(a[-1] ** 2) - return a[-1] ** 2 - - assert func([1, 2, 3]) == 9 - cmd.do('func 1 2 3') + assert a[1] == 2 + cmd.do("func 1 2 3") out, err = capsys.readouterr() - assert out == '9\n9\n' + assert out + err == '' - -def test_declare_command_list_float(capsys): +def test_declare_command_tuple(capsys): @cmd.declare_command - def func(a: List[float]): - print(a[-1]**2) - return a[-1]**2 - - assert func([1.1, 2.0, 3.0]) == 9.0 - cmd.do('func 1 2 3') + def func(a: Tuple[str, int]): + assert a == ("fooo", 42) + cmd.do("func fooo 42") out, err = capsys.readouterr() - assert out == '9.0\n9.0\n' + assert out + err == '' - -def test_declare_command_bool(capsys): +def test_declare_command_arg_docs(): @cmd.declare_command - def func(a: bool, b: bool): - assert a - assert not b + 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 - func(True, False) + 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] - cmd.do("func yes, no") +def test_declare_command_default(): + @cmd.declare_command + def func(a: str="sele"): + assert a == "a" + func("a") + cmd.do('func a') + + +def test_declare_command_enum(capsys): + class E(str, Enum): + A = "a" + B = "b" + C = "c" + + @cmd.declare_command + def func(e: E): + assert isinstance(e, E) + assert e == E.A + assert e == "a" + + cmd.do('func a') out, err = capsys.readouterr() - assert out == '' and err == '' \ No newline at end of file + assert out + err == '' \ No newline at end of file diff --git a/testing/tests/api/helping.py b/testing/tests/api/helping.py deleted file mode 100644 index d442015a6..000000000 --- a/testing/tests/api/helping.py +++ /dev/null @@ -1,50 +0,0 @@ -import sys -import unittest -from pymol import cmd, testing, stored - -try: - from io import StringIO - from unittest.mock import patch - mock_not_available = False -except ImportError: - mock_not_available = True - - -def func_with_indented_help(): - ''' - USAGE - - foo - - SEE ALSO - - https://github.com/schrodinger/pymol-open-source/issues/116 - ''' - - -cmd.extend('func_with_indented_help', func_with_indented_help) - - -@unittest.skipIf(mock_not_available, "unittest.mock not available") -class TestHelping(testing.PyMOLTestCase): - def testApi(self): - with patch('sys.stdout', new=StringIO()) as out: - cmd.api("color") - self.assertTrue('API: pymol.viewing.color' in out.getvalue()) - - def testHelp(self): - with patch('sys.stdout', new=StringIO()) as out: - cmd.help('color') - self.assertTrue('USAGE\n\n color color' in out.getvalue()) - - @testing.requires_version('2.5') - def testHelp_dedent(self): - with patch('sys.stdout', new=StringIO()) as out: - cmd.help('func_with_indented_help') - self.assertTrue('USAGE\n\n foo\n\nSEE' in out.getvalue()) - - @testing.requires_version('2.4') - @testing.requires('incentive') - def testHelpSetting(self): - out = cmd.help_setting('transparency') - self.assertTrue('controls surface transparency' in out)