Source code for deepmd.utils.plugin
"""Base of plugin systems."""
# copied from https://github.com/deepmodeling/dpdata/blob/a3e76d75de53f6076254de82d18605a010dc3b00/dpdata/plugin.py
from abc import ABCMeta
from typing import Callable
[docs]class Plugin:
"""A class to register and restore plugins.
Attributes
----------
plugins : Dict[str, object]
plugins
Examples
--------
>>> plugin = Plugin()
>>> @plugin.register("xx")
def xxx():
pass
>>> print(plugin.plugins['xx'])
"""
def __init__(self):
self.plugins = {}
def __add__(self, other) -> "Plugin":
self.plugins.update(other.plugins)
return self
[docs] def register(self, key : str) -> Callable[[object], object]:
"""Register a plugin.
Parameter
---------
key : str
key of the plugin
Returns
-------
Callable[[object], object]
decorator
"""
def decorator(object : object) -> object:
self.plugins[key] = object
return object
return decorator
[docs] def get_plugin(self, key) -> object:
"""Visit a plugin by key.
Parameters
----------
key : str
key of the plugin
Returns
-------
object
the plugin
"""
return self.plugins[key]
[docs]class VariantMeta:
def __call__(cls, *args, **kwargs):
"""Remove `type` and keys that starts with underline."""
obj = cls.__new__(cls, *args, **kwargs)
kwargs.pop('type', None)
to_pop = []
for kk in kwargs:
if kk[0] == '_':
to_pop.append(kk)
for kk in to_pop:
kwargs.pop(kk, None)
obj.__init__(*args, **kwargs)
return obj
[docs]class PluginVariant(metaclass=VariantABCMeta):
"""A class to remove `type` from input arguments."""
pass