pymol-open-source icon indicating copy to clipboard operation
pymol-open-source copied to clipboard

Renewed plugins

Open pslacerda opened this issue 4 months ago • 11 comments

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 avatar Mar 18 '24 08:03 pslacerda

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 avatar Mar 18 '24 21:03 pslacerda

Sample plugin command:

image

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

pslacerda avatar Mar 18 '24 21:03 pslacerda

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

pslacerda avatar Mar 20 '24 22:03 pslacerda

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

JarrettSJohnson avatar Mar 21 '24 18:03 JarrettSJohnson

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.

pslacerda avatar Mar 21 '24 19:03 pslacerda

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):
    ...

speleo3 avatar Mar 22 '24 18:03 speleo3

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

pslacerda avatar Mar 22 '24 18:03 pslacerda

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

pslacerda avatar Mar 22 '24 19:03 pslacerda

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

pslacerda avatar Mar 22 '24 20:03 pslacerda

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 avatar Mar 22 '24 21:03 pslacerda

ping

pslacerda avatar Apr 25 '24 21:04 pslacerda