mark_utils.py

"""Contains utility functions related to marker.

The utility functions are stored here to be separate from the plugin.

"""

from __future__ import annotations

from typing import TYPE_CHECKING
from typing import Any

from _pytask.models import CollectionMetadata
from _pytask.node_protocols import PTask
from _pytask.typing import TaskFunction
from _pytask.typing import attach_task_metadata
from _pytask.typing import is_task_decorator_target

if TYPE_CHECKING:
    from _pytask.mark import Mark


def get_all_marks(obj_or_task: Any | PTask) -> list[Mark]:
    """Get all marks from a callable or task."""
    if isinstance(obj_or_task, PTask):
        return obj_or_task.markers
    obj = obj_or_task
    return obj.pytask_meta.markers if isinstance(obj, TaskFunction) else []


def set_marks(obj_or_task: Any | PTask, marks: list[Mark]) -> Any | PTask:
    """Set marks on a callable or task."""
    if isinstance(obj_or_task, PTask):
        obj_or_task.markers = marks
    elif isinstance(obj_or_task, TaskFunction):
        obj_or_task.pytask_meta.markers = marks
    else:
        if not is_task_decorator_target(obj_or_task):
            msg = "Marks can only be stored on tasks or user-defined callables."
            raise TypeError(msg)
        attach_task_metadata(obj_or_task, CollectionMetadata(markers=marks))
    return obj_or_task


def get_marks(obj_or_task: Any | PTask, marker_name: str) -> list[Mark]:
    """Get marks from callable or task."""
    marks = get_all_marks(obj_or_task)
    return [mark for mark in marks if mark.name == marker_name]


def has_mark(obj_or_task: Any | PTask, marker_name: str) -> bool:
    """Test if callable or task has a certain mark."""
    marks = get_all_marks(obj_or_task)
    return any(mark.name == marker_name for mark in marks)


def remove_marks(
    obj_or_task: Any | PTask, marker_name: str
) -> tuple[Any | PTask, list[Mark]]:
    """Remove marks from callable or task."""
    marks = get_all_marks(obj_or_task)
    selected = [mark for mark in marks if mark.name == marker_name]
    others = [mark for mark in marks if mark.name != marker_name]
    obj_or_task = set_marks(obj_or_task, others)
    return obj_or_task, selected