schrodinger / pymol-open-source

Open-source foundation of the user-sponsored PyMOL molecular visualization system.
https://pymol.org/
Other
1.21k stars 284 forks source link

Renewed plugins #347

Open pslacerda opened 8 months ago

pslacerda commented 8 months ago

In my opinion, extended PyMOL commands should be both more friendly and strict.

This is a proof of concept:

from pathlib import Path

@declare_plugin
def f(a: int, b: Path='/tmp/'):
    return f'{a} {b}'

This is the messy code:

def declare_plugin(name, function=None, _self=cmd):
    if function is None:
        name, function = name.__name__, name

    if function.__code__.co_argcount != len(function.__annotations__):
        raise Exception("Messy annotations")

    def inner(*args, **kwargs):
        eval_args = []
        eval_kwargs = {}

        args = list(args[:])
        kwargs = kwargs.copy()

        anns = list(function.__annotations__.keys())
        funcs = list(function.__annotations__.values())

        for idx, arg in enumerate(args):
            key = anns[idx]
            func = funcs[idx]
            eval_args.append(func(arg))

        for kwarg_key, kwarg_arg in kwargs.items():
            key = kwarg_key
            arg = kwarg_arg
            func = function.__annotations__[key]

            func = function.__annotations__[key]
            eval_kwargs[kwarg_key] = func(kwarg_arg)

        print(eval_args, eval_kwargs)
        return function(*eval_args, **eval_kwargs)

    _self.keyword[name] = [inner, 0,0,',', parsing.STRICT]
    _self.kwhash.append(name)
    _self.help_sc.append(name)
    return inner
pslacerda commented 8 months ago

Here the updated code with autocomplete for Selection and Path.

from pathlib import Path

@declare_plugin
def f(a: int, b: Path='/tmp/'):
    return f'{a} {b}'

class Selection(str):
    pass

AUTO_ARGS = {}

def declare_plugin(name, function=None, _self=cmd):
    if function is None:
        name, function = name.__name__, name

    if function.__code__.co_argcount != len(function.__annotations__):
        raise Exception("Messy annotations")

    spec = inspect.getfullargspec(function)

    kwargs_ = {}
    args_ = spec.args[:]

    defaults = list(spec.defaults or [])

    while args_ and defaults:
        kwargs_[args_.pop(-1)] = defaults.pop(-1)

    for idx, (var, func) in enumerate(spec.annotations.items()):
        if (name, idx) not in AUTO_ARGS:
            auto = None
            if issubclass(func, Selection):
                auto = cmd.object_sc
            elif issubclass(func, Path):
                auto = lambda: cmd.Shortcut(glob("*"))
            if auto is not None:
                _self.cmd.auto_arg[idx][name] = [auto, "var", ""]

    def inner(*args, **kwargs):
        kwargs = {**dict(zip(args_, args)), **kwargs_, **kwargs}
        del kwargs["_self"]
        return function(**kwargs)

    name = function.__name__
    _self.keyword[name] = [inner, 0, 0, ",", parsing.STRICT]
    _self.kwhash.append(name)
    _self.help_sc.append(name)
    return inner
pslacerda commented 8 months ago

Sample plugin command:

image

More examples: https://gist.github.com/pslacerda/1e6a14cbe3fc7aeb2425e0a5d74ddc78

pslacerda commented 8 months ago

@JarrettSJohnson @speleo3, I'd like to work on this enhancement feature, can I?

JarrettSJohnson commented 8 months ago

This looks similar to the extend decorators we already have. Besides that, is this just enforcing type annotations on functions?

pslacerda commented 8 months ago

Yes, it is based the current approach with decorators.

It provides autocompletion for Selection args; does 'yes/no' conversion on Boolean; cast the args to the enforced types.

speleo3 commented 8 months ago

In principle, I like your idea and I think it would be cool if you work on that.

One thing: In my opinion, the type annotations must remain compatible with mypy. I'm not sure if that's the case with your current proposal.

By the way: There is also cmd.extendaa. It's not as elegant as your proposal, but still a convenient solution to add auto-completion.

@cmd.extendaa(
    [cmd.selection_sc, 'selection', ', '],
    [cmd.selection_sc, 'selection', ''],
)
def jaccard(sel1: str, sel2: str):
    ...
pslacerda commented 8 months ago

It is mostly mypy compatible and wrong only when the argument is with Bool/bool. image

pslacerda commented 8 months ago

Seems that to be both bool and coerce "yes"/"no" string is impossible because bool don't accept subclasses.

pslacerda commented 8 months ago

I was wrong and found a workaround for bool arguments. https://gist.github.com/pslacerda/1e6a14cbe3fc7aeb2425e0a5d74ddc78

pslacerda commented 8 months ago

There are more design analysis or can I already starts to implement on a branch?

I just implemented the support for Enums.

class FTMapEngine(StrEnum):
    FTMAP = "ftmap"
    ATLAS = "atlas"

@declare_plugin
def load_ftmap(
    filename: Path,
    group: str = "FTMap",
    max_length: int = 3,
    origin: FTMapEngine = FTMapEngine.FTMAP,
    max_combinations: int = 50,
    table: bool = True,
):
    ...

It's still missing support for other types like Any or Union[].

pslacerda commented 7 months ago

ping

JarrettSJohnson commented 1 month ago

Is this issue resolved now with #349 merged?

pslacerda commented 1 month ago

I think it's only partially resolved because it's missing error handling, testing and documentation.