collect_utils.py

"""Contains utility functions for :mod:`_pytask.collect`."""
from __future__ import annotations

import itertools
import uuid
import warnings
from pathlib import Path
from typing import Any
from typing import Callable
from typing import Generator
from typing import Iterable
from typing import TYPE_CHECKING

import attrs
from _pytask._inspect import get_annotations
from _pytask.exceptions import NodeNotCollectedError
from _pytask.mark_utils import has_mark
from _pytask.mark_utils import remove_marks
from _pytask.models import NodeInfo
from _pytask.node_protocols import PNode
from _pytask.nodes import PythonNode
from _pytask.shared import find_duplicates
from _pytask.task_utils import parse_keyword_arguments_from_signature_defaults
from _pytask.tree_util import PyTree
from _pytask.tree_util import tree_leaves
from _pytask.tree_util import tree_map
from _pytask.tree_util import tree_map_with_path
from _pytask.typing import no_default
from _pytask.typing import ProductType
from attrs import define
from attrs import field
from typing_extensions import Annotated
from typing_extensions import get_origin


if TYPE_CHECKING:
    from _pytask.session import Session


__all__ = [
    "depends_on",
    "parse_dependencies_from_task_function",
    "parse_products_from_task_function",
    "parse_nodes",
    "produces",
]


def depends_on(objects: PyTree[Any]) -> PyTree[Any]:
    """Specify dependencies for a task.

    Parameters
    ----------
    objects
        Can be any valid Python object or an iterable of any Python objects. To be
        valid, it must be parsed by some hook implementation for the
        :func:`_pytask.hookspecs.pytask_collect_node` entry-point.

    """
    return objects


def produces(objects: PyTree[Any]) -> PyTree[Any]:
    """Specify products of a task.

    Parameters
    ----------
    objects
        Can be any valid Python object or an iterable of any Python objects. To be
        valid, it must be parsed by some hook implementation for the
        :func:`_pytask.hookspecs.pytask_collect_node` entry-point.

    """
    return objects


def parse_nodes(  # noqa: PLR0913
    session: Session,
    task_path: Path | None,
    task_name: str,
    node_path: Path,
    obj: Any,
    parser: Callable[..., Any],
) -> Any:
    """Parse nodes from object."""
    arg_name = parser.__name__
    objects = _extract_nodes_from_function_markers(obj, parser)
    nodes = _convert_objects_to_node_dictionary(objects, arg_name)
    return tree_map(
        lambda x: _collect_decorator_node(
            session,
            node_path,
            task_name,
            NodeInfo(
                arg_name=arg_name,
                path=(),
                value=x,
                task_path=task_path,
                task_name=task_name,
            ),
        ),
        nodes,
    )


def _extract_nodes_from_function_markers(
    function: Callable[..., Any], parser: Callable[..., Any]
) -> Generator[Any, None, None]:
    """Extract nodes from a marker.

    The parser is a functions which is used to document the marker with the correct
    signature. Using the function as a parser for the ``args`` and ``kwargs`` of the
    marker provides the expected error message for misspecification.

    """
    marker_name = parser.__name__
    _, markers = remove_marks(function, marker_name)
    for marker in markers:
        parsed = parser(*marker.args, **marker.kwargs)
        yield parsed


def _convert_objects_to_node_dictionary(objects: Any, when: str) -> dict[Any, Any]:
    """Convert objects to node dictionary."""
    list_of_dicts = [_convert_to_dict(x) for x in objects]
    _check_that_names_are_not_used_multiple_times(list_of_dicts, when)
    return _merge_dictionaries(list_of_dicts)


@define(frozen=True)
class _Placeholder:
    """A placeholder to mark unspecified keys in dictionaries."""

    scalar: bool = field(default=False)
    id_: uuid.UUID = field(factory=uuid.uuid4)


def _convert_to_dict(x: Any, first_level: bool = True) -> Any | dict[Any, Any]:
    """Convert any object to a dictionary."""
    if isinstance(x, dict):
        return {k: _convert_to_dict(v, False) for k, v in x.items()}
    if isinstance(x, Iterable) and not isinstance(x, str):
        if first_level:
            return {
                _Placeholder(): _convert_to_dict(element, False)
                for i, element in enumerate(x)
            }
        return {i: _convert_to_dict(element, False) for i, element in enumerate(x)}
    if first_level:
        return {_Placeholder(scalar=True): x}
    return x


def _check_that_names_are_not_used_multiple_times(
    list_of_dicts: list[dict[Any, Any]], when: str
) -> None:
    """Check that names of nodes are not assigned multiple times.

    Tuples in the list have either one or two elements. The first element in the two
    element tuples is the name and cannot occur twice.

    """
    names_with_provisional_keys = list(
        itertools.chain.from_iterable(dict_.keys() for dict_ in list_of_dicts)
    )
    names = [x for x in names_with_provisional_keys if not isinstance(x, _Placeholder)]
    duplicated = find_duplicates(names)

    if duplicated:
        msg = f"'@pytask.mark.{when}' has nodes with the same name: {duplicated}"
        raise ValueError(msg)


def _union_of_dictionaries(dicts: list[dict[Any, Any]]) -> dict[Any, Any]:
    """Merge multiple dictionaries in one.

    Examples
    --------
    >>> a, b = {"a": 0}, {"b": 1}
    >>> _union_of_dictionaries([a, b])
    {'a': 0, 'b': 1}

    >>> a, b = {'a': 0}, {'a': 1}
    >>> _union_of_dictionaries([a, b])
    {'a': 1}

    """
    return dict(itertools.chain.from_iterable(dict_.items() for dict_ in dicts))


def _merge_dictionaries(list_of_dicts: list[dict[Any, Any]]) -> dict[Any, Any]:
    """Merge multiple dictionaries.

    The function does not perform a deep merge. It simply merges the dictionary based on
    the first level keys which are either unique names or placeholders. During the merge
    placeholders will be replaced by an incrementing integer.

    Examples
    --------
    >>> a, b = {"a": 0}, {"b": 1}
    >>> _merge_dictionaries([a, b])
    {'a': 0, 'b': 1}

    >>> a, b = {_Placeholder(): 0}, {_Placeholder(): 1}
    >>> _merge_dictionaries([a, b])
    {0: 0, 1: 1}

    """
    merged_dict = _union_of_dictionaries(list_of_dicts)

    if len(merged_dict) == 1 and isinstance(next(iter(merged_dict)), _Placeholder):
        placeholder, value = next(iter(merged_dict.items()))
        out = value if placeholder.scalar else {0: value}
    else:
        counter = itertools.count()
        out = {}
        for k, v in merged_dict.items():
            if isinstance(k, _Placeholder):
                while True:
                    possible_key = next(counter)
                    if possible_key not in merged_dict and possible_key not in out:
                        out[possible_key] = v
                        break
            else:
                out[k] = v

    return out


_ERROR_MULTIPLE_DEPENDENCY_DEFINITIONS = """The task uses multiple ways to define \
dependencies. Dependencies should be defined with either

- as default value for the function argument 'depends_on'.
- as '@pytask.task(kwargs={"depends_on": ...})'
- or with the deprecated '@pytask.mark.depends_on' and a 'depends_on' function argument.

Use only one of the two ways!

Hint: You do not need to use 'depends_on' as the argument name since pytask v0.4. \
Every function argument that is not a product is treated as a dependency. Read more \
about dependencies in the documentation: https://tinyurl.com/pytask-deps-prods.
"""


def parse_dependencies_from_task_function(
    session: Session, task_path: Path | None, task_name: str, node_path: Path, obj: Any
) -> dict[str, Any]:
    """Parse dependencies from task function."""
    has_depends_on_decorator = False
    has_depends_on_argument = False
    dependencies = {}

    if has_mark(obj, "depends_on"):
        has_depends_on_decorator = True
        nodes = parse_nodes(session, task_path, task_name, node_path, obj, depends_on)
        dependencies["depends_on"] = nodes

    task_kwargs = obj.pytask_meta.kwargs if hasattr(obj, "pytask_meta") else {}
    signature_defaults = parse_keyword_arguments_from_signature_defaults(obj)
    kwargs = {**signature_defaults, **task_kwargs}
    kwargs.pop("produces", None)

    # Parse products from task decorated with @task and that uses produces.
    if "depends_on" in kwargs:
        has_depends_on_argument = True
        dependencies["depends_on"] = tree_map(
            lambda x: _collect_decorator_node(
                session,
                node_path,
                task_name,
                NodeInfo(
                    arg_name="depends_on",
                    path=(),
                    value=x,
                    task_path=task_path,
                    task_name=task_name,
                ),
            ),
            kwargs["depends_on"],
        )

    if has_depends_on_decorator and has_depends_on_argument:
        raise NodeNotCollectedError(_ERROR_MULTIPLE_DEPENDENCY_DEFINITIONS)

    parameters_with_product_annot = _find_args_with_product_annotation(obj)
    parameters_with_node_annot = _find_args_with_node_annotation(obj)

    # Complete kwargs with node annotations, when no value is given by kwargs.
    for name in list(parameters_with_node_annot):
        if name not in kwargs:
            kwargs[name] = parameters_with_node_annot.pop(name)
        else:
            msg = (
                f"The value for the parameter {name!r} is defined twice in "
                "'@pytask.mark.task(kwargs=...)' and in the type annotation. Choose "
                "only one option."
            )
            raise ValueError(msg)

    for parameter_name, value in kwargs.items():
        if (
            parameter_name in parameters_with_product_annot
            or parameter_name == "return"
        ):
            continue

        if parameter_name == "depends_on":
            continue

        nodes = tree_map_with_path(
            lambda p, x: _collect_dependency(
                session,
                node_path,
                task_name,
                NodeInfo(
                    arg_name=parameter_name,  # noqa: B023
                    path=p,
                    value=x,
                    task_path=task_path,
                    task_name=task_name,
                ),
            ),
            value,
        )

        # If all nodes are python nodes, we simplify the parameter value and store it in
        # one node. If it is a node, we keep it.
        are_all_nodes_python_nodes_without_hash = all(
            isinstance(x, PythonNode) and not x.hash for x in tree_leaves(nodes)
        )
        if not isinstance(nodes, PNode) and are_all_nodes_python_nodes_without_hash:
            node_name = create_name_of_python_node(
                NodeInfo(
                    arg_name=parameter_name,
                    path=(),
                    value=value,
                    task_path=task_path,
                    task_name=task_name,
                )
            )
            dependencies[parameter_name] = PythonNode(value=value, name=node_name)
        else:
            dependencies[parameter_name] = nodes
    return dependencies


def _find_args_with_node_annotation(func: Callable[..., Any]) -> dict[str, PNode]:
    """Find args with node annotations."""
    annotations = get_annotations(func, eval_str=True)
    metas = {
        name: annotation.__metadata__
        for name, annotation in annotations.items()
        if get_origin(annotation) is Annotated
    }

    args_with_node_annotation = {}
    for name, meta in metas.items():
        annot = [i for i in meta if not isinstance(i, ProductType)]
        if len(annot) >= 2:  # noqa: PLR2004
            msg = (
                f"Parameter {name!r} has multiple node annotations although only one "
                f"is allowed. Annotations: {annot}"
            )
            raise ValueError(msg)
        if annot:
            args_with_node_annotation[name] = annot[0]

    return args_with_node_annotation


_ERROR_MULTIPLE_PRODUCT_DEFINITIONS = """The task uses multiple ways to define \
products. Products should be defined with either

- 'typing.Annotated[Path(...), Product]' (recommended)
- '@pytask.mark.task(kwargs={'produces': Path(...)})'
- as a default argument for 'produces': 'produces = Path(...)'
- '@pytask.mark.produces(Path(...))' (deprecated)

Read more about products in the documentation: https://tinyurl.com/pytask-deps-prods.
"""


def parse_products_from_task_function(
    session: Session, task_path: Path | None, task_name: str, node_path: Path, obj: Any
) -> dict[str, Any]:
    """Parse products from task function.

    Raises
    ------
    NodeNotCollectedError
        If multiple ways were used to specify products.

    """
    has_produces_decorator = False
    has_produces_argument = False
    has_annotation = False
    has_return = False
    has_task_decorator = False
    out = {}

    # Parse products from decorators.
    if has_mark(obj, "produces"):
        has_produces_decorator = True
        nodes = parse_nodes(session, task_path, task_name, node_path, obj, produces)
        out = {"produces": nodes}

    task_kwargs = obj.pytask_meta.kwargs if hasattr(obj, "pytask_meta") else {}
    signature_defaults = parse_keyword_arguments_from_signature_defaults(obj)
    kwargs = {**signature_defaults, **task_kwargs}

    parameters_with_product_annot = _find_args_with_product_annotation(obj)
    parameters_with_node_annot = _find_args_with_node_annotation(obj)

    # Parse products from task decorated with @task and that uses produces.
    if "produces" in kwargs:
        has_produces_argument = True
        collected_products = tree_map_with_path(
            lambda p, x: _collect_product(
                session,
                node_path,
                task_name,
                NodeInfo(
                    arg_name="produces",
                    path=p,
                    value=x,
                    task_path=task_path,
                    task_name=task_name,
                ),
                is_string_allowed=True,
            ),
            kwargs["produces"],
        )
        out = {"produces": collected_products}

    if parameters_with_product_annot:
        out = {}
        for parameter_name in parameters_with_product_annot:
            has_annotation = True
            if (
                parameter_name not in kwargs
                and parameter_name not in parameters_with_node_annot
            ):
                continue

            if (
                parameter_name in kwargs
                and parameter_name in parameters_with_node_annot
            ):
                msg = (
                    f"The value for the parameter {parameter_name!r} is defined twice "
                    "in '@pytask.mark.task(kwargs=...)' and in the type annotation. "
                    "Choose only one option."
                )
                raise ValueError(msg)

            value = kwargs.get(parameter_name) or parameters_with_node_annot.get(
                parameter_name
            )

            collected_products = tree_map_with_path(
                lambda p, x: _collect_product(
                    session,
                    node_path,
                    task_name,
                    NodeInfo(
                        arg_name=parameter_name,  # noqa: B023
                        path=p,
                        value=x,
                        task_path=task_path,
                        task_name=task_name,
                    ),
                    is_string_allowed=False,
                ),
                value,
            )
            out[parameter_name] = collected_products

    if "return" in parameters_with_node_annot:
        has_return = True
        collected_products = tree_map_with_path(
            lambda p, x: _collect_product(
                session,
                node_path,
                task_name,
                NodeInfo(
                    arg_name="return",
                    path=p,
                    value=x,
                    task_path=task_path,
                    task_name=task_name,
                ),
                is_string_allowed=False,
            ),
            parameters_with_node_annot["return"],
        )
        out = {"return": collected_products}

    task_produces = obj.pytask_meta.produces if hasattr(obj, "pytask_meta") else None
    if task_produces:
        has_task_decorator = True
        collected_products = tree_map_with_path(
            lambda p, x: _collect_product(
                session,
                node_path,
                task_name,
                NodeInfo(
                    arg_name="return",
                    path=p,
                    value=x,
                    task_path=task_path,
                    task_name=task_name,
                ),
                is_string_allowed=False,
            ),
            task_produces,
        )
        out = {"return": collected_products}

    if (
        sum(
            (
                has_produces_decorator,
                has_produces_argument,
                has_annotation,
                has_return,
                has_task_decorator,
            )
        )
        >= 2  # noqa: PLR2004
    ):
        raise NodeNotCollectedError(_ERROR_MULTIPLE_PRODUCT_DEFINITIONS)

    return out


def _find_args_with_product_annotation(func: Callable[..., Any]) -> list[str]:
    """Find args with product annotations."""
    annotations = get_annotations(func, eval_str=True)
    metas = {
        name: annotation.__metadata__
        for name, annotation in annotations.items()
        if get_origin(annotation) is Annotated
    }

    args_with_product_annot = []
    for name, meta in metas.items():
        has_product_annot = any(isinstance(i, ProductType) for i in meta)
        if has_product_annot:
            args_with_product_annot.append(name)

    return args_with_product_annot


_ERROR_WRONG_TYPE_DECORATOR = """'@pytask.mark.depends_on', '@pytask.mark.produces', \
and their function arguments can only accept values of type 'str' and 'pathlib.Path' \
or the same values nested in tuples, lists, and dictionaries. Here, {node} has type \
{node_type}.
"""


_WARNING_STRING_DEPRECATED = """Using strings to specify a {kind} is deprecated. Pass \
a 'pathlib.Path' instead with 'Path("{node}")'.
"""


def _collect_decorator_node(
    session: Session, path: Path, name: str, node_info: NodeInfo
) -> PNode:
    """Collect nodes for a task.

    Raises
    ------
    NodeNotCollectedError
        If the node could not collected.

    """
    node = node_info.value
    kind = {"depends_on": "dependency", "produces": "product"}.get(node_info.arg_name)

    if not isinstance(node, (str, Path)):
        raise NodeNotCollectedError(
            _ERROR_WRONG_TYPE_DECORATOR.format(node=node, node_type=type(node))
        )

    if isinstance(node, str):
        warnings.warn(
            _WARNING_STRING_DEPRECATED.format(kind=kind, node=node),
            category=FutureWarning,
            stacklevel=1,
        )
        node = Path(node)
        node_info = node_info._replace(value=node)

    collected_node = session.hook.pytask_collect_node(
        session=session, path=path, node_info=node_info
    )
    if collected_node is None:
        msg = f"{node!r} cannot be parsed as a {kind} for task {name!r} in {path!r}."
        raise NodeNotCollectedError(msg)

    return collected_node


def _collect_dependency(
    session: Session, path: Path, name: str, node_info: NodeInfo
) -> PNode:
    """Collect nodes for a task.

    Raises
    ------
    NodeNotCollectedError
        If the node could not collected.

    """
    node = node_info.value

    if isinstance(node, PythonNode) and node.value is no_default:
        # If a node is a dependency and its value is not set, the node is a product in
        # another task and the value will be set there. Thus, we wrap the original node
        # in another node to retrieve the value after it is set.
        new_node = attrs.evolve(node, value=node)
        node_info = node_info._replace(value=new_node)

    collected_node = session.hook.pytask_collect_node(
        session=session, path=path, node_info=node_info
    )
    if collected_node is None:
        msg = (
            f"{node!r} cannot be parsed as a dependency for task {name!r} in {path!r}."
        )
        raise NodeNotCollectedError(msg)
    return collected_node


def _collect_product(
    session: Session,
    path: Path,
    task_name: str,
    node_info: NodeInfo,
    is_string_allowed: bool = False,
) -> PNode:
    """Collect products for a task.

    Defining products with strings is only allowed when using the decorator. Parameter
    defaults can only be :class:`pathlib.Path`s.

    Raises
    ------
    NodeNotCollectedError
        If the node could not collected.

    """
    node = node_info.value
    # For historical reasons, task.kwargs is like the deco and supports str and Path.
    if not isinstance(node, (str, Path)) and is_string_allowed:
        msg = (
            f"`@pytask.mark.task(kwargs={{'produces': ...}}` can only accept values of "
            "type 'str' and 'pathlib.Path' or the same values nested in tuples, lists, "
            f"and dictionaries. Here, {node} has type {type(node)}."
        )
        raise ValueError(msg)

    # If we encounter a string and it is allowed, convert it to a path.
    if isinstance(node, str) and is_string_allowed:
        node = Path(node)
        node_info = node_info._replace(value=node)

    collected_node = session.hook.pytask_collect_node(
        session=session, path=path, node_info=node_info
    )

    if collected_node is None:
        msg = (
            f"{node!r} can't be parsed as a product for task {task_name!r} in {path!r}."
        )
        raise NodeNotCollectedError(msg)

    return collected_node


def create_name_of_python_node(node_info: NodeInfo) -> str:
    """Create name of PythonNode."""
    prefix = (
        node_info.task_path.as_posix() + "::" + node_info.task_name
        if node_info.task_path
        else node_info.task_name
    )
    node_name = prefix + "::" + node_info.arg_name
    if node_info.path:
        suffix = "-".join(map(str, node_info.path))
        node_name += "::" + suffix
    return node_name