Source code for dstk.hooks.tools

"""
This module provides the `Hook` class, which serves as a wrapper for wrapping
callable methods into a strict single-argument interface. It is designed to
facilitate the creation of modular pipelines where each processing step (a hook)
receives one primary piece of data while still allowing for optional configuration
via keyword arguments.

Core functionalities include:

* Wrapping any callable function into a `Hook` object to standardize how it
  is called within a larger model.
* Enforcing a single-argument constraint at runtime, ensuring that each hook
  focuses on one primary input.
* Separating the "main" argument from optional keyword arguments (defaults),
  allowing for cleaner configuration of complex text processing steps.
* Validating method signatures to ensure compatibility before execution.

This module is particularly useful for building extensible NLP pipelines where
different linguistic tasks (e.g., lemmatization, tagging, or filtering) need
to be chained together in a consistent manner.
"""

from typing import Any, Callable
from inspect import signature, Signature, Parameter
from collections.abc import ValuesView




class Hook:
    """
    Represents a callable hook that wraps a single-argument function.

    A Hook encapsulates a method that must accept exactly one argument, enabling
    modular processing steps or callbacks within models.

    :param name: A descriptive name for the hook.
    :type name: str
    :param method: A callable that takes exactly one argument.
    :type method: Callable[..., Any]
    """

    def __init__(self, method: Callable[..., Any]):
        """
        Initialize the Hook with a target callable.
        """

        self.method: Callable[..., Any] = method
        self.kwargs: dict[str, Any] = {}
        self.main_arg: str = ""

    def __call__(self, input_data: Any) -> Any:
        """
        Execute the wrapped method with the provided arguments after
        validating its signature.

        :param input_data: The data to pass as the primary argument.
        :type input_data: Any

        :return: The result of the wrapped method execution.
        :rtype: Any
        """

        self._check_args()

        if self.main_arg:
            self.kwargs.update({self.main_arg: input_data})
            return self.method(**self.kwargs)
        else:
            return self.method(input_data, **self.kwargs)



    def set_default_args(self, default_args: dict[str, Any]) -> None:
        """
        Update the dictionary of default keyword arguments for the
        wrapped method.

        :param default_args: A dictionary of argument names and values.
        :type default_args: dict[str, Any]
        """

        self.kwargs.update(default_args)




    def set_main_arg(self, main_arg: str) -> None:
        """
        Define the name of the primary argument for the wrapped method.

        :param main_arg: The name (string) of the primary argument.
        :type main_arg: str
        """

        args: int = len(main_arg.split())
        if args > 1:
            raise ValueError(
                f"Main argument be a single word. You provided {args} words."
            )

        self.main_arg = main_arg


    def _check_args(self) -> None:
        """
        Validate that the wrapped method's signature allows for exactly one argument.
        """
        sig: Signature = signature(self.method)
        params: ValuesView[Parameter] = sig.parameters.values()

        positional_params: list = [
            param for param in params if param.kind in {Parameter.POSITIONAL_ONLY}
        ]

        if len(positional_params) > 1:
            raise ValueError(
                "A hook method must accept exactly one positional argument"
            )

        if len(positional_params) > 0 and positional_params[0].name in self.kwargs:
            raise ValueError(
                "You cannot set the positional argument as a default, as it will be automatically passed when the workflow is excecuted."
            )