Source code for utopya.eval._plot_func_resolver
"""Implements a plot function resolver that takes model-specific information
into account."""
import logging
import os
import traceback
from types import ModuleType
from typing import Dict
import dantro._import_tools
import dantro.plot.utils
from ..model_registry import ModelInfoBundle
log = logging.getLogger(__name__)
# -----------------------------------------------------------------------------
[docs]class PlotFuncResolver(dantro.plot.utils.PlotFuncResolver):
"""A utopya-specific plot function resolver for
:py:class:`~utopya.eval.plotmanager.PlotManager` that takes information
from the model info bundle into account.
"""
BASE_PKG = "utopya.eval.plots"
"""Which package to use as base package for relative module imports"""
[docs] def __init__(
self,
*,
_model_info_bundle: ModelInfoBundle,
**kwargs,
):
"""Initializes the plot function resolver and additionally stores the
model info bundle.
"""
super().__init__(**kwargs)
self._model_info_bundle = _model_info_bundle
[docs] def _get_custom_module_paths(self) -> Dict[str, str]:
"""Aggregates a dict of module paths from which imports are attempted.
Uses model- or project-specific information, if available:
- A model's ``py_plots_dir``
- A project's ``py_plots_dir``
- A project's additional ``py_modules``
- The framework's additional ``py_modules``
- The framework's ``py_plots_dir``
"""
p = dict()
mib = self._model_info_bundle
if mib is None:
return p
if mib.paths.get("py_plots_dir"):
p["model's py_plots_dir"] = mib.paths["py_plots_dir"]
project = mib.project
if project:
if project.paths.py_plots_dir:
p["project's py_plots_dir"] = project.paths.py_plots_dir
if project.custom_py_modules:
for label, mod_path in project.custom_py_modules.items():
p[f"custom module '{label}'"] = mod_path
if project.framework_project is not None:
fw = project.framework_project
if fw.paths.py_plots_dir:
p["framework's py_plots_dir"] = fw.paths.py_plots_dir
if fw.custom_py_modules:
for label, mod_path in fw.custom_py_modules.items():
p[f"framework's custom module '{label}'"] = mod_path
return p
[docs] def _get_module_via_import(self, *, module: str, **kwargs) -> ModuleType:
"""Extends the parent method by making the custom modules available if
the regular import failed.
"""
try:
return super()._get_module_via_import(module=module, **kwargs)
except ModuleNotFoundError as err:
# Are there additional modules that are to be searched for imports?
custom_module_paths = self._get_custom_module_paths()
if not custom_module_paths:
log.note(
"No custom module paths available to import '%s' from.",
module,
)
raise
log.debug(
"Module '%s' could not be imported with a default sys.path, "
"but custom plot modules are available. Attempting to "
"import it with %d additional path(s) being available.",
module,
len(custom_module_paths),
)
# Go over the specified custom paths and try to import them, gathering
# detailed error information if that fails
errors = dict()
for key, mod_path in custom_module_paths.items():
# In order to be able to import modules at the given path, the
# sys.path needs to include the _parent_ directory of this path.
parent_dir = os.path.dirname(mod_path)
# Enter two context managers, taking care to return both sys.path
# and sys.modules back to the same state as they were before their
# invocation.
# The latter context manager is crucial because module imports lead
# to a cache entry even if a subsequent attempt to import a part of
# the module string was the cause of an error (which makes sense).
# Example: a failing `model_plots.foo` import would still lead to a
# cache entry of the `model_plots` module; however, attempting to
# then import `model_plots.bar` will make the lookup _only_ in the
# cached module. As we want several import attempts here, the cache
# is not desired.
add_sys_path = dantro._import_tools.added_sys_path(parent_dir)
tmp_sys_modules = dantro._import_tools.temporary_sys_modules(
reset_only_on_fail=True
)
with add_sys_path, tmp_sys_modules:
try:
mod = super()._get_module_via_import(
module=module, **kwargs
)
except ModuleNotFoundError as err:
_tb = err.__traceback__
errors[parent_dir] = dict(
err=err,
tb=_tb,
tb_lines=traceback.format_tb(_tb),
)
else:
log.debug(
"Found module '%s' after having added custom plot "
"module path labelled '%s' (%s) to the sys.path.",
mod,
key,
mod_path,
)
return mod
# All imports failed. Inform extensively about errors to help debugging
err_info = "\n".join(
f"-- Error at custom plot module path {p} : {e['err']}\n\n"
" Abbreviated traceback:\n"
f"{e['tb_lines'][0]} ...\n"
f"{e['tb_lines'][-1]}"
for p, e in errors.items()
)
raise ModuleNotFoundError(
f"Could not import module '{module}'! It was found neither among "
"the installed packages nor among the custom plot modules.\n"
"\n"
"The following errors were encountered at the respective custom "
"plot module search paths:\n\n"
f"{err_info}\n"
"NOTE: This error can have two reasons:\n"
f" (1) the '{module}' module does not exist in the specified "
" search location.\n"
" (2) during import of the plot module you specified, an "
"_unrelated_ ModuleNotFoundError occurred somewhere inside _your_ "
"code.\n"
"To debug, check the error messages and tracebacks above to find "
"out which of the two is preventing module import."
)
