Source code for dantro.plot.utils.plot_func

"""Implements utilities that revolve around the plotting function which is then
invoked by the :ref:`plot creators <plot_creators>`:

- a decorator to declare a function as a plot function
- the tools to resolve a plotting function from a module or file
"""

import importlib
import importlib.util
import logging
import os
import warnings
from typing import Callable, Sequence, Union

from ..._import_tools import (
    import_module_from_file as _import_module_from_file,
)

log = logging.getLogger(__name__)

# -----------------------------------------------------------------------------


[docs]class PlotFuncResolver: """Takes care of resolving a plot function""" BASE_PKG: str = "dantro.plot.funcs" """The default module string to use for relative module imports, where this module becomes the base package. Evaluated in :py:meth:`.__init__`. """
[docs] def __init__( self, *, base_module_file_dir: str = None, base_pkg: str = None ): """Set up the plot function resolver. Args: base_module_file_dir (str, optional): If given, ``module_file`` arguments to :py:meth:`.resolve` that are relative paths will be seen relative to this directory. Needs to be an absolute directory path and supports ``~`` expansion. base_pkg (str, optional): If given, use this base package instead for relative module imports instead of :py:attr:`.BASE_PKG`. Raises: ValueError: If ``base_module_file_dir`` was not absolute FileNotFoundError: If ``base_module_file_dir`` is missing or not a directory. """ if base_module_file_dir: bmfd = os.path.expanduser(base_module_file_dir) if not os.path.isabs(bmfd): raise ValueError( "Argument `base_module_file_dir` needs to be " f"an absolute path, was not! Got: {bmfd}" ) elif not os.path.exists(bmfd) or not os.path.isdir(bmfd): raise FileNotFoundError( "Argument `base_module_file_dir` does not " "exists or does not point to a directory!" ) self.base_module_file_dir = base_module_file_dir self.base_pkg = base_pkg if base_pkg else self.BASE_PKG
[docs] def resolve( self, *, plot_func: Union[str, Callable], module: str = None, module_file: str = None, ) -> Callable: """Resolve and return the plot function callable Args: plot_func (Union[str, Callable]): The name or module string of the plot function as it can be imported from ``module``. If this is a callable will directly return that callable. module (str): If ``plot_func`` was the name of the plot function, this needs to be the name of the module to import that name from. module_file (str): Path to the file to load and look for the ``plot_func`` in. If ``base_module_file_dir`` is given during initialization, this can also be a path relative to that directory. Returns: Callable: The resolved plot function Raises: TypeError: On bad argument types """ if callable(plot_func): log.debug("Received plotting function: %s", str(plot_func)) return self._attach_attributes(plot_func) elif not isinstance(plot_func, str): raise TypeError( "Argument `plot_func` needs to be a string or a " f"callable, was {type(plot_func)} with value '{plot_func}'." ) # else: need to resolve the module and find the callable in it # For less confusing variable names, do some renaming plot_func_modstr = plot_func del plot_func # First resolve the module, either from file or via import if module_file: mod = self._get_module_from_file( module_file, base_module_file_dir=self.base_module_file_dir ) elif isinstance(module, str): mod = self._get_module_via_import( module=module, base_pkg=self.base_pkg ) else: raise TypeError( "Could not import a module, because neither argument " "`module_file` was given nor did argument `module` have the " f"correct type (needs to be string but was {type(module)} " f"with value '{module}')." ) # plot_func could be something like "A.B.C.d"; go along the segments to # allow for more versatile plot function retrieval attr_names = plot_func_modstr.split(".") for attr_name in attr_names[:-1]: mod = getattr(mod, attr_name) # This is now the last module. Get the actual function plot_func = getattr(mod, attr_names[-1]) log.debug("Resolved plotting function: %s", str(plot_func)) # Decorate with some attributes, then return the result return self._attach_attributes( plot_func, module=mod, plot_func_modstr=plot_func_modstr )
[docs] def _get_module_from_file(self, path: str, *, base_module_file_dir: str): """Returns the module corresponding to the file at the given ``path``. This uses :py:func:`~dantro._import_tools.import_module_from_file` to carry out the import. """ try: return _import_module_from_file( path, base_dir=base_module_file_dir ) except ValueError as err: raise ValueError( "Need to specify `base_module_file_dir` during initialization " "to use relative paths for `module_file` argument!" ) from err
[docs] def _get_module_via_import(self, *, module: str, base_pkg: str): """Returns the module via import. Imports ``module`` via importlib, allowing relative imports from the package defined as base package. """ return importlib.import_module(module, package=base_pkg)
[docs] def _attach_attributes( self, plot_func: Callable, /, *, module=None, plot_func_modstr: str = None, ) -> Callable: """Attaches some informational attributes to the plot function.""" plot_func.modstr = plot_func_modstr # TODO attach a `name` here, while the information is available return plot_func
# -----------------------------------------------------------------------------
[docs]class is_plot_func: """This is a decorator class declaring the decorated function as a plotting function to use with :py:class:`~dantro.plot.creators.base.BasePlotCreator` or derived creators. .. note:: This decorator has a set of specializations that make sense only when using a specific creator type! For example, the ``helper``-related arguments are only used by :py:class:`~dantro.plot.creators.pyplot.PyPlotCreator` and are ignored without warning otherwise. """
[docs] def __init__( self, *, creator: str = None, creator_type: type = None, creator_name: str = None, use_dag: bool = None, required_dag_tags: Sequence[str] = None, compute_only_required_dag_tags: bool = True, pass_dag_object_along: bool = False, unpack_dag_results: bool = False, use_helper: bool = None, helper_defaults: Union[dict, str] = None, supports_animation=False, add_attributes: dict = None, ): """Initialize the decorator. .. note:: Some arguments are only evaluated when using a certain creator type, e.g. :py:class:`~dantro.plot.creators.pyplot.PyPlotCreator`. Args: creator (str, optional): The creator to use; needs to be registered with the PlotManager under this name. creator_type (type, optional): The type of plot creator to use. This argument is DEPRECATED, use ``creator`` instead. creator_name (str, optional): The name of the plot creator to use. This argument is DEPRECATED, use ``creator`` instead. use_dag (bool, optional): Whether to use the data transformation framework. required_dag_tags (Sequence[str], optional): The DAG tags that are required by the plot function. compute_only_required_dag_tags (bool, optional): Whether to compute only those DAG tags that are specified as required by the plot function. This is ignored if no required DAG tags were given and can be overwritten by the ``compute_only`` argument. pass_dag_object_along (bool, optional): Whether to pass on the DAG object to the plot function unpack_dag_results (bool, optional): Whether to unpack the results of the DAG computation directly into the plot function instead of passing it as a dictionary. use_helper (bool, optional): Whether to use the :py:class:`~dantro.plot.plot_helper.PlotHelper` with this plot. Needs :py:class:`~dantro.plot.creators.pyplot.PyPlotCreator`. If None, will default to True for supported creators and False otherwise. helper_defaults (Union[dict, str], optional): Default configurations for helpers; these are automatically considered to be enabled. If not dict-like, will assume this is an absolute path (supporting ``~`` expansion) to a YAML file and will load the dict-like configuration from there. Needs :py:class:`~dantro.plot.creators.pyplot.PyPlotCreator`. supports_animation (bool, optional): Whether the plot function supports animation. Needs :py:class:`~dantro.plot.creators.pyplot.PyPlotCreator`. add_attributes (dict, optional): Additional attributes to add to the plot function. Raises: ValueError: If ``helper_defaults`` was a string but not an absolute path. """ from ..._yaml import load_yml if helper_defaults and not isinstance(helper_defaults, dict): # Interpret as absolute path to yaml file fpath = os.path.expanduser(helper_defaults) if not os.path.isabs(fpath): raise ValueError( "`helper_defaults` string argument was a " f"relative path: {fpath}, but needs to be either a " "dict or an absolute path (~ allowed)." ) log.debug("Loading helper defaults from file %s ...", fpath) helper_defaults = load_yml(fpath) # Evaluate the creator argument if creator: if creator_name or creator_type: raise ValueError( "Cannot pass both `creator` and `creator_name` or " "`creator_type`!" ) else: if creator_name and creator_type: raise ValueError( "Cannot pass both of the deprecated decorator arguments " "`creator_name` and `creator_type`. Use `creator` instead." ) _warn_msg = ( "The `{}` argument is deprecated! Use `creator` instead." ) if creator_name: warnings.warn( _warn_msg.format("creator_name"), DeprecationWarning ) creator = creator_name elif creator_type: warnings.warn( _warn_msg.format("creator_type"), DeprecationWarning ) creator = creator_type # Gather those attributes that are to be set as function attributes self.pf_attrs = dict( is_plot_func=True, creator=creator, use_dag=use_dag, required_dag_tags=required_dag_tags, compute_only_required_dag_tags=compute_only_required_dag_tags, pass_dag_object_along=pass_dag_object_along, use_helper=use_helper, helper_defaults=helper_defaults, supports_animation=supports_animation, **(add_attributes if add_attributes else {}), )
[docs] def __call__(self, func: Callable): """If there are decorator arguments, __call__() is only called once, as part of the decoration process and expects as only argument the function to be decorated. """ # Do not actually wrap the function call, but add attributes to it for k, v in self.pf_attrs.items(): setattr(func, k, v) # Return the function, now with attributes set return func