Source code for genno.core.describe
from collections.abc import Hashable
from functools import partial
from itertools import chain
from textwrap import shorten
from typing import Any, Mapping
import dask.core
import xarray as xr
from .key import Key
#: Default maximum length for outputs from :func:`describe_recursive`.
MAX_ITEM_LENGTH = 160
[docs]
def describe_recursive(graph, comp, depth=0, seen=None):
"""Recursive helper for :meth:`.describe`.
Parameters
----------
graph :
A dask graph.
comp :
A dask computation.
depth : int
Recursion depth. Used for indentation.
seen : set
Keys that have already been described. Used to avoid
double-printing.
"""
comp = comp if isinstance(comp, tuple) else (comp,)
seen = set() if seen is None else seen
indent = (" " * 2 * (depth - 1)) + ("- " if depth > 0 else "")
# Strings for arguments
result = []
for arg in comp:
# Don't fully reprint keys and their ancestors that have been seen
if isinstance(arg, Hashable) and arg in seen:
if depth > 0:
# Don't print top-level items that have been seen
result.append(f"{indent}'{arg}' (above)")
continue
elif isinstance(arg, (str, Key)) and arg in graph:
# key that exists in the graph → recurse
item = "'{}':\n{}".format(
arg, describe_recursive(graph, graph[arg], depth + 1, seen)
)
elif is_list_of_keys(arg, graph):
# list → collection of items
item = "list of:\n{}".format(
describe_recursive(graph, tuple(arg), depth + 1, seen)
)
else:
# Anything else: use a readable string representation
item = label(arg)
try:
seen.add(arg)
except TypeError:
pass # `arg` is unhashable, e.g. a list
result.append(indent + item)
# Combine items
return ("\n" if depth > 0 else "\n\n").join(result)
[docs]
def is_list_of_keys(arg: Any, graph: Mapping) -> bool:
"""Identify a task which is a list of other keys."""
return (
isinstance(arg, list)
and len(arg) > 0
and isinstance(arg[0], Hashable)
and arg[0] in graph
)
[docs]
def label(arg, max_length=MAX_ITEM_LENGTH) -> str:
"""Return a label for `arg`.
The label depends on the type of `arg`:
- :class:`.xarray.DataArray`: the first line of the string representation.
- :func:`functools.partial` object: a less-verbose version that omits None
arguments.
- Item protected with :func:`.dask.core.quote`: its literal value.
- A callable, e.g. a function: its name.
- Anything else: its :class:`str` representation.
In all cases, the string is no longer than `max_length`.
"""
# Convert various types of arguments to string
if isinstance(arg, xr.DataArray):
# DataArray → just the first line of the string representation
return str(arg).split("\n")[0]
elif isinstance(arg, partial):
# functools.partial → less verbose format
fn_args = ", ".join(
chain(
map(repr, arg.args),
filter(
None,
map(
lambda kw: "" if kw[1] is None else f"{kw[0]}={kw[1]}",
arg.keywords.items(),
),
),
["..."],
)
)
fn_repr = getattr(arg.func, "__name__", repr(arg.func))
return shorten(f"{fn_repr}({fn_args})", max_length)
elif isinstance(arg, dask.core.literal):
# Item protected with dask.core.quote()
return shorten(str(arg.data), max_length)
elif callable(arg):
if arg.__class__.__name__ == "builtin_function_or_method":
return repr(arg).replace("function ", "")
else:
return getattr(arg, "__name__", str(arg))
else:
return shorten(str(arg), max_length)