From c68717c31b4bde1b1bb6ac4239dde5c9ee237af5 Mon Sep 17 00:00:00 2001 From: Pedro Sousa Lacerda Date: Sat, 29 Nov 2025 17:29:36 -0300 Subject: [PATCH 1/6] Remove unsuitable tests --- testing/tests/api/helping.py | 50 ------------------------------------ 1 file changed, 50 deletions(-) delete mode 100644 testing/tests/api/helping.py 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) From cf40a4f83a5358b08ff8b9c6eefd9408ab442e6d Mon Sep 17 00:00:00 2001 From: Pedro Sousa Lacerda Date: Sat, 29 Nov 2025 17:31:15 -0300 Subject: [PATCH 2/6] Small refactoration of a testing file --- testing/tests/api/test_commanding.py | 113 ++++++++++++++------------- 1 file changed, 57 insertions(+), 56 deletions(-) diff --git a/testing/tests/api/test_commanding.py b/testing/tests/api/test_commanding.py index 3bb5ee526..448ae4df5 100644 --- a/testing/tests/api/test_commanding.py +++ b/testing/tests/api/test_commanding.py @@ -7,93 +7,93 @@ def test_docstring(): @cmd.new_command - def func1(): + def func(): """docstring""" - assert func1.__doc__ == "docstring" + assert func.__doc__ == "docstring" -@cmd.new_command -def func2(a: bool, b: bool): - assert a - assert not b def test_bool(capsys): - cmd.do("func2 yes, 0") + @cmd.new_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 == '' -@cmd.new_command -def func3( - 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" - + def test_generic(capsys): - cmd.do("func3 nullable_point=1 2 3, my_foo=11.0") + @cmd.new_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 == '' -@cmd.new_command -def func4(dirname: Path = Path('.')): - assert dirname.exists() - def test_path(capsys): - cmd.do('func4 ..') - cmd.do('func4') + @cmd.new_command + def func(dirname: Path = Path('.')): + assert dirname.exists() + cmd.do('func ..') + cmd.do('func') out, err = capsys.readouterr() assert out + err == '' -@cmd.new_command -def func5(old_style: Any): - assert old_style is RuntimeError -func5(RuntimeError) + @mark.skip("This function does not works as expected") def test_any(capsys): - - cmd.do("func5 RuntimeError") + @cmd.new_command + def func(old_style: Any): + assert old_style is RuntimeError + func(RuntimeError) + cmd.do("func RuntimeError") out, err = capsys.readouterr() assert 'AssertionError' not in out+err -@cmd.new_command -def func6(a: List): - assert a[1] == "2" - -@cmd.new_command -def func7(a: List[int]): - assert a[1] == 2 - def test_list(capsys): - cmd.do("func6 1 2 3") + @cmd.new_command + def func(a: List): + assert a[1] == "2" + + cmd.do("func 1 2 3") out, err = capsys.readouterr() assert out + err == '' - cmd.do("func7 1 2 3") + @cmd.new_command + def func(a: List[int]): + assert a[1] == 2 + + cmd.do("func 1 2 3") out, err = capsys.readouterr() assert out + err == '' -@cmd.new_command -def func8(a: Tuple[str, int]): - assert a == ("fooo", 42) - def test_tuple(capsys): - cmd.do("func8 fooo 42") + @cmd.new_command + def func(a: Tuple[str, int]): + assert a == ("fooo", 42) + + cmd.do("func fooo 42") out, err = capsys.readouterr() assert out + err == '' -@cmd.new_command -def func10(a: str="sele"): - assert a == "sele" def test_default(capsys): - cmd.do('func10') + @cmd.new_command + def func(a: str="sele"): + assert a == "sele" + + cmd.do('func') out, err = capsys.readouterr() assert out + err == '' @@ -106,9 +106,10 @@ def test_str_enum(capsys): class E(StrEnum): A = "a" @cmd.new_command - def func11(e: E): + def func(e: E): assert e == E.A assert isinstance(e, E) - cmd.do('func11 a') + cmd.do('func a') out, err = capsys.readouterr() - assert out + err == '' \ No newline at end of file + assert out + err == '' + From 2417417267c5ba68159d68a803304d9790b71b79 Mon Sep 17 00:00:00 2001 From: Pedro Sousa Lacerda Date: Sat, 29 Nov 2025 17:34:25 -0300 Subject: [PATCH 3/6] Add support to quiet argument on new_command --- modules/pymol/commanding.py | 17 ++++++++++------- testing/tests/api/test_commanding.py | 7 +++++++ 2 files changed, 17 insertions(+), 7 deletions(-) diff --git a/modules/pymol/commanding.py b/modules/pymol/commanding.py index 6b4e3d074..3bc8ef247 100644 --- a/modules/pymol/commanding.py +++ b/modules/pymol/commanding.py @@ -689,18 +689,21 @@ def inner(*args, **kwargs): # It was called from command line or pml script, so parse arguments if caller == _parser_filename: kwargs = {**kwargs, **dict(zip(args2_, args))} + # special _self argument 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 + # special 'quiet' argument + if var == 'quiet' and isinstance(value, int): + new_kwargs[var] = bool(value) + else: + new_kwargs[var] = _into_types(type, value) + final_kwargs = { + **kwargs_, + **new_kwargs + } return function(**final_kwargs) # It was called from Python, so pass the arguments as is diff --git a/testing/tests/api/test_commanding.py b/testing/tests/api/test_commanding.py index 448ae4df5..fc4f81076 100644 --- a/testing/tests/api/test_commanding.py +++ b/testing/tests/api/test_commanding.py @@ -113,3 +113,10 @@ def func(e: E): out, err = capsys.readouterr() assert out + err == '' +def test_quiet(capsys): + @cmd.new_command + def func(quiet: bool=True): + assert not quiet + cmd.do('func') + out, err = capsys.readouterr() + assert out + err == '' \ No newline at end of file From ddb5be6f835f73ec9abeb6fe1bac43d0cd180e2e Mon Sep 17 00:00:00 2001 From: Pedro Sousa Lacerda Date: Sat, 29 Nov 2025 17:42:32 -0300 Subject: [PATCH 4/6] Fix error on typed lists --- modules/pymol/commanding.py | 10 ++++++---- testing/tests/api/test_commanding.py | 26 +++++++++++++++++--------- 2 files changed, 23 insertions(+), 13 deletions(-) diff --git a/modules/pymol/commanding.py b/modules/pymol/commanding.py index 3bc8ef247..99b9df5f7 100644 --- a/modules/pymol/commanding.py +++ b/modules/pymol/commanding.py @@ -641,11 +641,13 @@ def _into_types(type, value): elif issubclass(list, origin): args = get_args(type) - if len(args) > 0: + if len(args) == 1: f = args[0] - else: - f = lambda x: x - return [f(i) for i in shlex.split(value)] + return [ + _into_types(f, a) + for a in shlex.split(value) + ] + return shlex.split(value) elif issubclass(type, Enum): if value in type: diff --git a/testing/tests/api/test_commanding.py b/testing/tests/api/test_commanding.py index fc4f81076..c8e081fc2 100644 --- a/testing/tests/api/test_commanding.py +++ b/testing/tests/api/test_commanding.py @@ -1,9 +1,8 @@ -from pytest import mark -from pymol import cmd import sys +from pytest import mark from typing import List, Union, Any, Tuple from pathlib import Path - +from pymol import cmd def test_docstring(): @cmd.new_command @@ -40,6 +39,7 @@ def func( out, err = capsys.readouterr() assert out + err == '' + def test_path(capsys): @cmd.new_command def func(dirname: Path = Path('.')): @@ -50,7 +50,6 @@ def func(dirname: Path = Path('.')): assert out + err == '' - @mark.skip("This function does not works as expected") def test_any(capsys): @cmd.new_command @@ -61,11 +60,11 @@ def func(old_style: Any): out, err = capsys.readouterr() assert 'AssertionError' not in out+err + def test_list(capsys): @cmd.new_command def func(a: List): assert a[1] == "2" - cmd.do("func 1 2 3") out, err = capsys.readouterr() assert out + err == '' @@ -73,17 +72,24 @@ def func(a: List): @cmd.new_command def func(a: List[int]): assert a[1] == 2 - cmd.do("func 1 2 3") out, err = capsys.readouterr() assert out + err == '' + @cmd.new_command + def func(a: List[bool]): + assert a.pop(0) == False + assert a.pop(0) == True + cmd.do("func 0 yes") + out, err = capsys.readouterr() + assert out + err == '' + + def test_tuple(capsys): @cmd.new_command def func(a: Tuple[str, int]): - assert a == ("fooo", 42) - - cmd.do("func fooo 42") + assert a == ("fooo a", 42) + cmd.do("func 'fooo a' 42") out, err = capsys.readouterr() assert out + err == '' @@ -97,6 +103,7 @@ def func(a: str="sele"): out, err = capsys.readouterr() assert out + err == '' + @mark.skipif( sys.version_info < (3, 11), reason="Requires StrEnum of Python 3.11+" @@ -113,6 +120,7 @@ def func(e: E): out, err = capsys.readouterr() assert out + err == '' + def test_quiet(capsys): @cmd.new_command def func(quiet: bool=True): From 8bfa31c399545d6098d90a0711b2ee4b78a6c666 Mon Sep 17 00:00:00 2001 From: Pedro Sousa Lacerda Date: Sat, 29 Nov 2025 19:46:43 -0300 Subject: [PATCH 5/6] 20x speedup with ChatGPT tip --- modules/pymol/commanding.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/modules/pymol/commanding.py b/modules/pymol/commanding.py index 99b9df5f7..9c1b5ba9e 100644 --- a/modules/pymol/commanding.py +++ b/modules/pymol/commanding.py @@ -687,7 +687,7 @@ def new_command(name, function=None, _self=cmd): # Inner function that will be callable every time the command is executed @wraps(function) def inner(*args, **kwargs): - caller = traceback.extract_stack(limit=2)[0].filename + caller = sys._getframe(1).f_code.co_filename # It was called from command line or pml script, so parse arguments if caller == _parser_filename: kwargs = {**kwargs, **dict(zip(args2_, args))} From 2c0175d581350e9cbaf9262655950caa6ffe548d Mon Sep 17 00:00:00 2001 From: Pedro Sousa Lacerda Date: Sun, 30 Nov 2025 00:53:42 -0300 Subject: [PATCH 6/6] Better error messages and add support to Enum and UnionType --- modules/pymol/commanding.py | 151 +++++++++++++++++++-------- testing/tests/api/test_commanding.py | 31 +++++- 2 files changed, 133 insertions(+), 49 deletions(-) diff --git a/modules/pymol/commanding.py b/modules/pymol/commanding.py index 9c1b5ba9e..f4d68cbd1 100644 --- a/modules/pymol/commanding.py +++ b/modules/pymol/commanding.py @@ -24,16 +24,16 @@ from io import FileIO as file import inspect - import glob import shlex - import tokenize import builtins - from io import BytesIO from enum import Enum + if sys.version_info >= (3, 11): + from enum import StrEnum 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 + from typing import get_args, Union, Any, get_origin + from types import UnionType import re import os @@ -602,63 +602,121 @@ def get_state_list(states_str): output = get_state_list(states) states_list = sorted(set(map(int, output))) return _cmd.delete_states(_self._COb, name, states_list) + + + class ArgumentParsingError(ValueError): + "Error on argument parsing." - def _into_types(type, value): - if repr(type) == 'typing.Any': + def __init__(self, arg_name, message): + message = dedent(message).strip() + if arg_name: + s = f"Failed at parsing '{arg_name}'. {message}" + else: + s = message + super().__init__(s) + + + def _into_types(var, type, value): + + # Untyped string + if type == Any: return value + + # Boolean flags elif type is bool: if isinstance(value, bool): return value - if value.lower() in ["yes", "1", "true", "on", "y"]: + trues = ["yes", "1", "true", "on", "y"] + falses = ["no", "0", "false", "off", "n"] + if value.lower() in trues: return True - elif value.lower() in ["no", "0", "false", "off", "n"]: + elif value.lower() in falses: return False else: - raise pymol.CmdException(f"Invalid boolean value: {value}") + raise ArgumentParsingError( + var, + f"Can't parse {value!r} as bool." + f" Supported true values are {', '.join(trues)}." + f" Supported false values are {', '.join(falses)}." + ) - 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) + # Types from typing module + elif origin := get_origin(type): + + if origin in {Union, UnionType}: + funcs = get_args(type) + for func in funcs: + try: + return _into_types(None, func, value) + except: + continue + raise ArgumentParsingError( + var, + f"Can't parse {value!r} into {type}." + f" The parser tried each union type and none was suitable." + ) - elif origin == Union: - args = get_args(type) - found = False - for i, arg in enumerate(args): + elif issubclass(origin, tuple): + funcs = get_args(type) + if funcs: + values = shlex.split(value) + if len(funcs) > 0 and len(funcs) != len(values): + raise ArgumentParsingError( + var, + f"Can't parse {value!r} into {type}." + f" The number of tuple arguments are incorrect." + ) try: - found = True - return _into_types(arg, value) + return tuple(_into_types(None, f, v) for f, v in zip(funcs, values)) except: - found = False - if not found: - raise pymol.CmdException(f"Union was not able to cast {value}") - - elif issubclass(list, origin): - args = get_args(type) - if len(args) == 1: - f = args[0] - return [ - _into_types(f, a) - for a in shlex.split(value) - ] + raise ArgumentParsingError( + var, + f"Can't parse {value!r} into {type}." + f" One or more tuple values are of incorrect types." + ) + else: + return tuple(shlex.split(value)) + + elif issubclass(origin, list): + funcs = get_args(type) + if len(funcs) == 1: + func = funcs[0] + return [_into_types(None, func, a) for a in shlex.split(value)] return shlex.split(value) - elif issubclass(type, Enum): - if value in type: + elif sys.version_info >= (3, 11) and issubclass(type, StrEnum): + try: return type(value) - else: - raise pymol.CmdException(f"Invalid value for enum {type.__name__}: {value}") + except: + names = [e.value for e in list(type)] + raise ArgumentParsingError( + var, + f"Invalid value for {type.__name__}." + f" Accepted values are {', '.join(names)}." + ) + + # Specific types must go before other generic types + # isinstance(type, builtins.type) comes after + elif issubclass(type, Enum): + value = type.__members__.get(value) + if value is None: + raise ArgumentParsingError( + var, + f"Invalid value for {type.__name__}." + f" Accepted values are {', '.join(type.__members__)}." + ) + return value - elif isinstance(type, str): - return str(value) - - raise pymol.CmdException(f"Unsupported argument type annotation {type}") + # Generic types must accept str as single argument to __init__(s) + elif isinstance(type, builtins.type): + try: + return type(value) + except Exception as exc: + raise ArgumentParsingError( + var, + f"Invalid value {value!r} for custom type {type.__name__}." + f" The type must accept str as the solo argument to __init__(s)." + ) from exc def new_command(name, function=None, _self=cmd): @@ -701,7 +759,7 @@ def inner(*args, **kwargs): if var == 'quiet' and isinstance(value, int): new_kwargs[var] = bool(value) else: - new_kwargs[var] = _into_types(type, value) + new_kwargs[var] = _into_types(var, type, value) final_kwargs = { **kwargs_, **new_kwargs @@ -719,6 +777,7 @@ def inner(*args, **kwargs): inner.func = inner.__wrapped__ return inner + def extend(name, function=None, _self=cmd): ''' diff --git a/testing/tests/api/test_commanding.py b/testing/tests/api/test_commanding.py index c8e081fc2..a74a114b9 100644 --- a/testing/tests/api/test_commanding.py +++ b/testing/tests/api/test_commanding.py @@ -1,8 +1,11 @@ import sys from pytest import mark -from typing import List, Union, Any, Tuple +from typing import List, Union, Any, Tuple, Optional from pathlib import Path + from pymol import cmd +from pymol.commanding import ArgumentParsingError + def test_docstring(): @cmd.new_command @@ -26,7 +29,8 @@ def test_generic(capsys): def func( nullable_point: Tuple[float, float, float], my_var: Union[int, float] = 10, - my_foo: Union[int, float] = 10.0, + my_foo: int | float = 10.0, + null_ptr: Optional[bool] = None, extended_calculation: bool = True, old_style: Any = "Old behavior" ): @@ -34,6 +38,7 @@ def func( assert extended_calculation assert isinstance(my_var, int) assert isinstance(my_foo, float) + assert null_ptr is None assert old_style == "Old behavior" cmd.do("func nullable_point=1 2 3, my_foo=11.0") out, err = capsys.readouterr() @@ -104,6 +109,20 @@ def func(a: str="sele"): assert out + err == '' +def test_enum(capsys): + from enum import Enum + class E(Enum): + A = 1 + B = 2 + @cmd.new_command + def func(e: E): + assert e == E.A + assert isinstance(e, E) + cmd.do('func A') + out, err = capsys.readouterr() + assert out + err == '' + + @mark.skipif( sys.version_info < (3, 11), reason="Requires StrEnum of Python 3.11+" @@ -112,6 +131,7 @@ def test_str_enum(capsys): from enum import StrEnum class E(StrEnum): A = "a" + B = "b" @cmd.new_command def func(e: E): assert e == E.A @@ -127,4 +147,9 @@ def func(quiet: bool=True): assert not quiet cmd.do('func') out, err = capsys.readouterr() - assert out + err == '' \ No newline at end of file + assert out + err == '' + + +def test_argument_error(): + err = ArgumentParsingError('my_var', "Short error message.") + assert str(err) == "Failed at parsing 'my_var'. Short error message." \ No newline at end of file