Enhance handling of optional imports and programs for pytask dag. (#155)

Author: tobiasraabe

docs/source/changes.rst

@@ -20,6 +20,8 @@ all releases are available on `PyPI <https://pypi.org/project/pytask>`_ and
   ``n_entries_in_table`` in the configuration file.
 - :gh:`152` makes the duration of the execution readable by humans by separating it into
   days, hours, minutes and seconds.
+- :gh:`155` implements functions to check for optional packages and programs and raises
+  errors for requirements to draw the DAG earlier.
 
 
 0.1.1 - 2021-08-25

src/_pytask/compat.py

@@ -0,0 +1,136 @@
+"""This module contains functions to assess compatibility and optional dependencies."""
+import importlib
+import shutil
+import sys
+import types
+import warnings
+from typing import Optional
+
+from packaging.version import parse as parse_version
+
+
+_MINIMUM_VERSIONS = {}
+"""Dict[str, str]: A mapping from packages to their minimum versions."""
+
+
+_IMPORT_TO_PACKAGE_NAME = {}
+"""Dict[str, str]: A mapping from import name to package name (on PyPI) for packages
+where these two names are different."""
+
+
+def _get_version(module: types.ModuleType) -> str:
+    version = getattr(module, "__version__", None)
+    if version is None:
+        raise ImportError(f"Can't determine version for {module.__name__}")
+    return version
+
+
+def import_optional_dependency(
+    name: str,
+    extra: str = "",
+    errors: str = "raise",
+    min_version: Optional[str] = None,
+    caller: str = "pytask",
+) -> Optional[types.ModuleType]:
+    """Import an optional dependency.
+
+    By default, if a dependency is missing an ImportError with a nice message will be
+    raised. If a dependency is present, but too old, we raise.
+
+    Parameters
+    ----------
+    name : str
+        The module name.
+    extra : str
+        Additional text to include in the ImportError message.
+    errors : str {'raise', 'warn', 'ignore'}
+        What to do when a dependency is not found or its version is too old.
+
+        * raise : Raise an ImportError
+        * warn : Only applicable when a module's version is to old. Warns that the
+          version is too old and returns None
+        * ignore: If the module is not installed, return None, otherwise, return the
+          module, even if the version is too old. It's expected that users validate the
+          version locally when using ``errors="ignore"`` (see. ``io/html.py``)
+    min_version : str, default None
+        Specify a minimum version that is different from the global pandas minimum
+        version required.
+    caller : str, default "pytask"
+        The caller of the function.
+
+    Returns
+    -------
+    maybe_module : Optional[ModuleType]
+        The imported module, when found and the version is correct. None is returned
+        when the package is not found and `errors` is False, or when the package's
+        version is too old and `errors` is ``'warn'``.
+
+    """
+    if errors not in ("warn", "raise", "ignore"):
+        raise ValueError("'errors' must be one of 'warn', 'raise' or 'ignore'.")
+
+    package_name = _IMPORT_TO_PACKAGE_NAME.get(name)
+    install_name = package_name if package_name is not None else name
+
+    if extra and not extra.endswith(" "):
+        extra += " "
+    msg = (
+        f"{caller} requires the optional dependency '{install_name}'. {extra}"
+        f"Use pip or conda to install '{install_name}'."
+    )
+    try:
+        module = importlib.import_module(name)
+    except ImportError:
+        if errors == "raise":
+            raise ImportError(msg) from None
+        else:
+            return None
+
+    # Handle submodules: if we have submodule, grab parent module from sys.modules
+    parent = name.split(".")[0]
+    if parent != name:
+        install_name = parent
+        module_to_get = sys.modules[install_name]
+    else:
+        module_to_get = module
+    minimum_version = (
+        min_version if min_version is not None else _MINIMUM_VERSIONS.get(parent)
+    )
+    if minimum_version:
+        version = _get_version(module_to_get)
+        if parse_version(version) < parse_version(minimum_version):
+            msg = (
+                f"{caller} requires version '{minimum_version}' or newer of "
+                f"'{parent}' (version '{version}' currently installed)."
+            )
+            if errors == "warn":
+                warnings.warn(msg, UserWarning)
+                return None
+            elif errors == "raise":
+                raise ImportError(msg)
+
+    return module
+
+
+def check_for_optional_program(
+    name: str,
+    extra: str = "",
+    errors: str = "raise",
+    caller: str = "pytask",
+) -> Optional[bool]:
+    if errors not in ("warn", "raise", "ignore"):
+        raise ValueError(
+            f"'errors' must be one of 'warn', 'raise' or 'ignore' and not '{errors}'."
+        )
+
+    msg = f"{caller} requires the optional program '{name}'. {extra}"
+
+    program_exists = shutil.which(name) is not None
+
+    if not program_exists:
+        if errors == "raise":
+            raise RuntimeError(msg)
+        elif errors == "warn":
+            warnings.warn(msg, UserWarning)
+
+    return program_exists

src/_pytask/graph.py

@@ -1,11 +1,13 @@
 """This file contains the command and code for drawing the DAG."""
-import shutil
+import sys
 from pathlib import Path
 from typing import Any
 from typing import Dict
 
 import click
 import networkx as nx
+from _pytask.compat import check_for_optional_program
+from _pytask.compat import import_optional_dependency
 from _pytask.config import hookimpl
 from _pytask.console import console
 from _pytask.dag import descending_tasks
@@ -18,6 +20,8 @@
 from _pytask.session import Session
 from _pytask.shared import get_first_non_none_value
 from _pytask.shared import reduce_names_of_multiple_nodes
+from _pytask.traceback import remove_internal_traceback_frames_from_exc_info
+from rich.traceback import Traceback
 
 
 @hookimpl(tryfirst=True)
@@ -61,9 +65,50 @@ def pytask_parse_config(config, config_from_cli, config_from_file):
 @click.option("-o", "--output-path", type=str, default=None, help=_HELP_TEXT_OUTPUT)
 def dag(**config_from_cli):
     """Create a visualization of the project's DAG."""
-    session = _create_session(config_from_cli)
-    dag = _refine_dag(session)
-    _write_graph(dag, session.config["output_path"], session.config["layout"])
+    try:
+        pm = get_plugin_manager()
+        from _pytask import cli
+
+        pm.register(cli)
+        pm.hook.pytask_add_hooks(pm=pm)
+
+        config = pm.hook.pytask_configure(pm=pm, config_from_cli=config_from_cli)
+
+        session = Session.from_config(config)
+
+    except (ConfigurationError, Exception):
+        console.print_exception()
+        session = Session({}, None)
+        session.exit_code = ExitCode.CONFIGURATION_FAILED
+
+    else:
+        try:
+            session.hook.pytask_log_session_header(session=session)
+            import_optional_dependency("pydot")
+            check_for_optional_program(
+                session.config["layout"],
+                extra="The layout program is part of the graphviz package which you "
+                "can install with conda.",
+            )
+            session.hook.pytask_collect(session=session)
+            session.hook.pytask_resolve_dependencies(session=session)
+            dag = _refine_dag(session)
+            _write_graph(dag, session.config["output_path"], session.config["layout"])
+
+        except CollectionError:
+            session.exit_code = ExitCode.COLLECTION_FAILED
+
+        except ResolvingDependenciesError:
+            session.exit_code = ExitCode.RESOLVING_DEPENDENCIES_FAILED
+
+        except Exception:
+            session.exit_code = ExitCode.FAILED
+            exc_info = remove_internal_traceback_frames_from_exc_info(sys.exc_info())
+            console.print()
+            console.print(Traceback.from_exception(*exc_info))
+            console.rule(style=ColorCode.FAILED)
+
+    sys.exit(session.exit_code)
 
 
 def build_dag(config_from_cli: Dict[str, Any]) -> "pydot.Dot":  # noqa: F821
@@ -87,9 +132,40 @@ def build_dag(config_from_cli: Dict[str, Any]) -> "pydot.Dot":  # noqa: F821
         A preprocessed graph which can be customized and exported.
 
     """
-    session = _create_session(config_from_cli)
-    dag = _refine_dag(session)
-    return dag
+    try:
+        pm = get_plugin_manager()
+        from _pytask import cli
+
+        pm.register(cli)
+        pm.hook.pytask_add_hooks(pm=pm)
+
+        config = pm.hook.pytask_configure(pm=pm, config_from_cli=config_from_cli)
+
+        session = Session.from_config(config)
+
+    except (ConfigurationError, Exception):
+        console.print_exception()
+        session = Session({}, None)
+        session.exit_code = ExitCode.CONFIGURATION_FAILED
+
+    else:
+        try:
+            session.hook.pytask_log_session_header(session=session)
+            import_optional_dependency("pydot")
+            check_for_optional_program(
+                session.config["layout"],
+                extra="The layout program is part of the graphviz package which you "
+                "can install with conda.",
+            )
+            session.hook.pytask_collect(session=session)
+            session.hook.pytask_resolve_dependencies(session=session)
+            dag = _refine_dag(session)
+
+        except Exception:
+            raise
+
+        else:
+            return dag
 
 
 def _refine_dag(session):
@@ -122,6 +198,8 @@ def _create_session(config_from_cli: Dict[str, Any]) -> nx.DiGraph:
     else:
         try:
             session.hook.pytask_log_session_header(session=session)
+            import_optional_dependency("pydot")
+            check_for_optional_program(session.config["layout"])
             session.hook.pytask_collect(session=session)
             session.hook.pytask_resolve_dependencies(session=session)
 
@@ -186,21 +264,6 @@ def _escape_node_names_with_colons(dag: nx.DiGraph):
 
 
 def _write_graph(dag: nx.DiGraph, path: Path, layout: str) -> None:
-    try:
-        import pydot  # noqa: F401
-    except ImportError:
-        raise ImportError(
-            "To visualize the project's DAG you need to install pydot which is "
-            "available with pip and conda. For example, use 'conda install -c "
-            "conda-forge pydot'."
-        ) from None
-    if shutil.which(layout) is None:
-        raise RuntimeError(
-            f"The layout program '{layout}' could not be found on your PATH. Please, "
-            "install graphviz. For example, use 'conda install -c conda-forge "
-            "graphivz'."
-        )
-
     path.parent.mkdir(exist_ok=True, parents=True)
     graph = nx.nx_pydot.to_pydot(dag)
     graph.write(path, prog=layout, format=path.suffix[1:])