Source code for dantro.data_ops.db_tools

"""Tools that help to monitor and manipulate the operations database"""

import logging
from difflib import get_close_matches as _get_close_matches
from typing import Callable, Dict, Sequence, Union

from ..exceptions import *
from import make_columns as _make_columns
from .db import _OPERATIONS

log = logging.getLogger(__name__)

# -----------------------------------------------------------------------------
# TODO Let operations registry use ObjectRegistry as basis

[docs]def register_operation( func: Callable, name: str = None, *, skip_existing: bool = False, overwrite_existing: bool = False, _ops: dict = None, ) -> None: """Adds an entry to the shared operations registry. Args: func (Callable): The callable that is to be registered as operation. name (str, optional): The name of the operation. If not given (and the callable not being a lambda), will use the function name instead. skip_existing (bool, optional): Whether to skip registration if the operation name is already registered. This suppresses the ValueError raised on existing operation name. overwrite_existing (bool, optional): Whether to overwrite a potentially already existing operation of the same name. If given, this takes precedence over ``skip_existing``. _ops (dict, optional): The operations database object to use; if None, uses the :ref:`dantro operations database <data_ops_available>`. Raises: TypeError: On invalid name or non-callable for the func argument ValueError: On already existing operation name and no skipping or overwriting enabled. Also if no ``name`` was given but the given callable is a lambda (which only has ``<lambda>`` as name). """ if _ops is None: _ops = _OPERATIONS if not callable(func): raise TypeError( f"The given {func} for operation '{name}' is not callable!" ) if name is None: name = func.__name__ if name == (lambda: 0).__name__: raise ValueError( "Could not automatically deduce an operation name because the " "given callable appears to be a lambda! Explicitly specify an " "operation name using the `name` argument." ) elif not isinstance(name, str): raise TypeError( f"Operation name need be a string, was {type(name)} with " f"value {name}!" ) if name in _ops and not overwrite_existing: if skip_existing: log.debug( "Operation '%s' is already registered and will not be " "registered again. Choose a different name or unset the " "`skip_existing` flag.", name, ) return raise ValueError( f"Operation name '{name}' already exists! Not overwriting. " "Set the `overwrite_existing` flag to force overwriting or choose " "a different `name`. If no name was given, consider specifying it " "explicitly instead of letting it be deduced." ) _ops[name] = func log.debug("Registered operation '%s'.", name)
[docs]def is_operation( arg: Union[str, Callable] = None, /, *, _ops: dict = None, _reg_func: Callable = None, **kws, ): """Decorator for registering functions with the operations database. As an alternative to :py:func:`.register_operation`, this decorator can be used to register a function with the operations database right where its defined: .. testcode:: from dantro.data_ops import is_operation # Operation name deduced from function name @is_operation def my_operation(data, *args): # ... do stuff here ... return data # Custom operation name @is_operation("do_something") def my_operation_with_a_custom_name(foo, bar): pass # Overwriting an operation of the same name @is_operation("do_something", overwrite_existing=True) def actually_do_something(spam, fish): pass .. testcode:: :hide: from dantro.data_ops.db import _OPERATIONS assert "my_operation" in _OPERATIONS assert "do_something" in _OPERATIONS See :ref:`register_data_ops` for general information. For instructions on how to overwrite this decorator with a custom one, see :ref:`customize_db_tools`. Args: arg (Union[str, Callable], optional): The name that should be used in the operation registry. If not given, will use the name of the decorated function instead. If a callable, this refers to the ``@is_operation`` call syntax and will use that as a function. _ops (dict, optional): The operations database to use. If not given, uses the :ref:`dantro operations database <data_ops_available>`. _reg_func (Callable, optional): If given, uses that callable for registration, which should have the same signature as :py:func:`.register_operation`. If None, uses dantro's registration function, :py:func:`.register_operation`. **kws: Passed to :py:func:`.register_operation` or a potentially given custom ``_reg_func``. """ if _reg_func is None: _reg_func = register_operation def wrapper(func): # Need some nonlocal names, see PEP3104 nonlocal arg, _reg_func name = arg if callable(name): # Invocation via @is_operation -- without parentheses name = None elif name is None: # Invocation via @is_operation() -- without arguments pass _reg_func(func, name=name, _ops=_ops, **kws) return func # Allow both invocation styles if callable(arg): # @is_operation return wrapper(arg) # @is_operation(*args, **kwargs) return wrapper
[docs]def available_operations( *, match: str = None, n: int = 5, _ops: dict = None ) -> Sequence[str]: """Returns all available operation names or a fuzzy-matched subset of them. Also see :ref:`data_ops_available` for an overview. Args: match (str, optional): If given, fuzzy-matches the names and only returns close matches to this name. n (int, optional): Number of close matches to return. Passed on to :py:func:`difflib.get_close_matches` _ops (dict, optional): The operations database object to use; if None, uses the :ref:`dantro operations database <data_ops_available>`. Returns: Sequence[str]: All available operation names or the matched subset. The sequence is sorted alphabetically. """ if _ops is None: _ops = _OPERATIONS if match is None: return _ops.keys() # Use fuzzy matching to return close matches return _get_close_matches(match, _ops.keys(), n=n)
[docs]def get_operation( op_name: str, *, _ops: Dict[str, Callable] = None ) -> Callable: """Retrieve the operation's callable Args: op_name (str): Name of the operation _ops (Dict[str, Callable], optional): The operations database object to use; if None, uses the :ref:`dantro operations database <data_ops_available>`. Raises: BadOperationName: Upon invalid operation name """ try: return _ops[op_name] except KeyError as err: # Find some close matches to make operation discovery easier _close_matches = available_operations(match=op_name, n=8, _ops=_ops) _did_you_mean = ( f" Did you mean: {', '.join(_close_matches)} ?" if _close_matches else "" ) _available = _make_columns(available_operations(_ops=_ops)) raise BadOperationName( f"No operation '{op_name}' registered!{_did_you_mean} " f"\nAvailable operations:\n{_available}If you need to register " "a new operation, use dantro.utils.register_operation." ) from err