Skip to content

Commit f3bc48e

Browse files
tobiasraabetobiasraabe
authored
Enhance handling of optional imports and programs for pytask dag. (#155)
1 parent 8975284 commit f3bc48e

File tree

6 files changed

+494
-27
lines changed

6 files changed

+494
-27
lines changed

‎docs/source/changes.rst

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

2426

2527
0.1.1 - 2021-08-25

‎src/_pytask/compat.py

+136
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,136 @@
1+
"""This module contains functions to assess compatibility and optional dependencies."""
2+
import importlib
3+
import shutil
4+
import sys
5+
import types
6+
import warnings
7+
from typing import Optional
8+
9+
from packaging.version import parse as parse_version
10+
11+
12+
_MINIMUM_VERSIONS = {}
13+
"""Dict[str, str]: A mapping from packages to their minimum versions."""
14+
15+
16+
_IMPORT_TO_PACKAGE_NAME = {}
17+
"""Dict[str, str]: A mapping from import name to package name (on PyPI) for packages
18+
where these two names are different."""
19+
20+
21+
def _get_version(module: types.ModuleType) -> str:
22+
version = getattr(module, "__version__", None)
23+
if version is None:
24+
raise ImportError(f"Can't determine version for {module.__name__}")
25+
return version
26+
27+
28+
def import_optional_dependency(
29+
name: str,
30+
extra: str = "",
31+
errors: str = "raise",
32+
min_version: Optional[str] = None,
33+
caller: str = "pytask",
34+
) -> Optional[types.ModuleType]:
35+
"""Import an optional dependency.
36+
37+
By default, if a dependency is missing an ImportError with a nice message will be
38+
raised. If a dependency is present, but too old, we raise.
39+
40+
Parameters
41+
----------
42+
name : str
43+
The module name.
44+
extra : str
45+
Additional text to include in the ImportError message.
46+
errors : str {'raise', 'warn', 'ignore'}
47+
What to do when a dependency is not found or its version is too old.
48+
49+
* raise : Raise an ImportError
50+
* warn : Only applicable when a module's version is to old. Warns that the
51+
version is too old and returns None
52+
* ignore: If the module is not installed, return None, otherwise, return the
53+
module, even if the version is too old. It's expected that users validate the
54+
version locally when using ``errors="ignore"`` (see. ``io/html.py``)
55+
min_version : str, default None
56+
Specify a minimum version that is different from the global pandas minimum
57+
version required.
58+
caller : str, default "pytask"
59+
The caller of the function.
60+
61+
Returns
62+
-------
63+
maybe_module : Optional[ModuleType]
64+
The imported module, when found and the version is correct. None is returned
65+
when the package is not found and `errors` is False, or when the package's
66+
version is too old and `errors` is ``'warn'``.
67+
68+
"""
69+
if errors not in ("warn", "raise", "ignore"):
70+
raise ValueError("'errors' must be one of 'warn', 'raise' or 'ignore'.")
71+
72+
package_name = _IMPORT_TO_PACKAGE_NAME.get(name)
73+
install_name = package_name if package_name is not None else name
74+
75+
if extra and not extra.endswith(" "):
76+
extra += " "
77+
msg = (
78+
f"{caller} requires the optional dependency '{install_name}'. {extra}"
79+
f"Use pip or conda to install '{install_name}'."
80+
)
81+
try:
82+
module = importlib.import_module(name)
83+
except ImportError:
84+
if errors == "raise":
85+
raise ImportError(msg) from None
86+
else:
87+
return None
88+
89+
# Handle submodules: if we have submodule, grab parent module from sys.modules
90+
parent = name.split(".")[0]
91+
if parent != name:
92+
install_name = parent
93+
module_to_get = sys.modules[install_name]
94+
else:
95+
module_to_get = module
96+
minimum_version = (
97+
min_version if min_version is not None else _MINIMUM_VERSIONS.get(parent)
98+
)
99+
if minimum_version:
100+
version = _get_version(module_to_get)
101+
if parse_version(version) < parse_version(minimum_version):
102+
msg = (
103+
f"{caller} requires version '{minimum_version}' or newer of "
104+
f"'{parent}' (version '{version}' currently installed)."
105+
)
106+
if errors == "warn":
107+
warnings.warn(msg, UserWarning)
108+
return None
109+
elif errors == "raise":
110+
raise ImportError(msg)
111+
112+
return module
113+
114+
115+
def check_for_optional_program(
116+
name: str,
117+
extra: str = "",
118+
errors: str = "raise",
119+
caller: str = "pytask",
120+
) -> Optional[bool]:
121+
if errors not in ("warn", "raise", "ignore"):
122+
raise ValueError(
123+
f"'errors' must be one of 'warn', 'raise' or 'ignore' and not '{errors}'."
124+
)
125+
126+
msg = f"{caller} requires the optional program '{name}'. {extra}"
127+
128+
program_exists = shutil.which(name) is not None
129+
130+
if not program_exists:
131+
if errors == "raise":
132+
raise RuntimeError(msg)
133+
elif errors == "warn":
134+
warnings.warn(msg, UserWarning)
135+
136+
return program_exists

‎src/_pytask/graph.py

+85-22
Original file line numberDiff line numberDiff line change
@@ -1,11 +1,13 @@
11
"""This file contains the command and code for drawing the DAG."""
2-
import shutil
2+
import sys
33
from pathlib import Path
44
from typing import Any
55
from typing import Dict
66

77
import click
88
import networkx as nx
9+
from _pytask.compat import check_for_optional_program
10+
from _pytask.compat import import_optional_dependency
911
from _pytask.config import hookimpl
1012
from _pytask.console import console
1113
from _pytask.dag import descending_tasks
@@ -18,6 +20,8 @@
1820
from _pytask.session import Session
1921
from _pytask.shared import get_first_non_none_value
2022
from _pytask.shared import reduce_names_of_multiple_nodes
23+
from _pytask.traceback import remove_internal_traceback_frames_from_exc_info
24+
from rich.traceback import Traceback
2125

2226

2327
@hookimpl(tryfirst=True)
@@ -61,9 +65,50 @@ def pytask_parse_config(config, config_from_cli, config_from_file):
6165
@click.option("-o", "--output-path", type=str, default=None, help=_HELP_TEXT_OUTPUT)
6266
def dag(**config_from_cli):
6367
"""Create a visualization of the project's DAG."""
64-
session = _create_session(config_from_cli)
65-
dag = _refine_dag(session)
66-
_write_graph(dag, session.config["output_path"], session.config["layout"])
68+
try:
69+
pm = get_plugin_manager()
70+
from _pytask import cli
71+
72+
pm.register(cli)
73+
pm.hook.pytask_add_hooks(pm=pm)
74+
75+
config = pm.hook.pytask_configure(pm=pm, config_from_cli=config_from_cli)
76+
77+
session = Session.from_config(config)
78+
79+
except (ConfigurationError, Exception):
80+
console.print_exception()
81+
session = Session({}, None)
82+
session.exit_code = ExitCode.CONFIGURATION_FAILED
83+
84+
else:
85+
try:
86+
session.hook.pytask_log_session_header(session=session)
87+
import_optional_dependency("pydot")
88+
check_for_optional_program(
89+
session.config["layout"],
90+
extra="The layout program is part of the graphviz package which you "
91+
"can install with conda.",
92+
)
93+
session.hook.pytask_collect(session=session)
94+
session.hook.pytask_resolve_dependencies(session=session)
95+
dag = _refine_dag(session)
96+
_write_graph(dag, session.config["output_path"], session.config["layout"])
97+
98+
except CollectionError:
99+
session.exit_code = ExitCode.COLLECTION_FAILED
100+
101+
except ResolvingDependenciesError:
102+
session.exit_code = ExitCode.RESOLVING_DEPENDENCIES_FAILED
103+
104+
except Exception:
105+
session.exit_code = ExitCode.FAILED
106+
exc_info = remove_internal_traceback_frames_from_exc_info(sys.exc_info())
107+
console.print()
108+
console.print(Traceback.from_exception(*exc_info))
109+
console.rule(style=ColorCode.FAILED)
110+
111+
sys.exit(session.exit_code)
67112

68113

69114
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
87132
A preprocessed graph which can be customized and exported.
88133
89134
"""
90-
session = _create_session(config_from_cli)
91-
dag = _refine_dag(session)
92-
return dag
135+
try:
136+
pm = get_plugin_manager()
137+
from _pytask import cli
138+
139+
pm.register(cli)
140+
pm.hook.pytask_add_hooks(pm=pm)
141+
142+
config = pm.hook.pytask_configure(pm=pm, config_from_cli=config_from_cli)
143+
144+
session = Session.from_config(config)
145+
146+
except (ConfigurationError, Exception):
147+
console.print_exception()
148+
session = Session({}, None)
149+
session.exit_code = ExitCode.CONFIGURATION_FAILED
150+
151+
else:
152+
try:
153+
session.hook.pytask_log_session_header(session=session)
154+
import_optional_dependency("pydot")
155+
check_for_optional_program(
156+
session.config["layout"],
157+
extra="The layout program is part of the graphviz package which you "
158+
"can install with conda.",
159+
)
160+
session.hook.pytask_collect(session=session)
161+
session.hook.pytask_resolve_dependencies(session=session)
162+
dag = _refine_dag(session)
163+
164+
except Exception:
165+
raise
166+
167+
else:
168+
return dag
93169

94170

95171
def _refine_dag(session):
@@ -122,6 +198,8 @@ def _create_session(config_from_cli: Dict[str, Any]) -> nx.DiGraph:
122198
else:
123199
try:
124200
session.hook.pytask_log_session_header(session=session)
201+
import_optional_dependency("pydot")
202+
check_for_optional_program(session.config["layout"])
125203
session.hook.pytask_collect(session=session)
126204
session.hook.pytask_resolve_dependencies(session=session)
127205

@@ -186,21 +264,6 @@ def _escape_node_names_with_colons(dag: nx.DiGraph):
186264

187265

188266
def _write_graph(dag: nx.DiGraph, path: Path, layout: str) -> None:
189-
try:
190-
import pydot # noqa: F401
191-
except ImportError:
192-
raise ImportError(
193-
"To visualize the project's DAG you need to install pydot which is "
194-
"available with pip and conda. For example, use 'conda install -c "
195-
"conda-forge pydot'."
196-
) from None
197-
if shutil.which(layout) is None:
198-
raise RuntimeError(
199-
f"The layout program '{layout}' could not be found on your PATH. Please, "
200-
"install graphviz. For example, use 'conda install -c conda-forge "
201-
"graphivz'."
202-
)
203-
204267
path.parent.mkdir(exist_ok=True, parents=True)
205268
graph = nx.nx_pydot.to_pydot(dag)
206269
graph.write(path, prog=layout, format=path.suffix[1:])

0 commit comments

Comments
 (0)