validmind 2.5.24__py3-none-any.whl → 2.6.7__py3-none-any.whl

validmind/tests/decorator.py

@@ -2,216 +2,28 @@
 # See the LICENSE file in the root of this repository for details.
 # SPDX-License-Identifier: AGPL-3.0 AND ValidMind Commercial
 
-"""Decorators for creating and registering metrics with the ValidMind framework."""
-
-# TODO: as we move entirely to a functional approach a lot of this logic
-# should be moved into the __init__ to replace the old class-based stuff
+"""Decorators for creating and registering tests with the ValidMind Library."""
 
 import inspect
 import os
-from typing import Any, Dict, List, Tuple, Union
-from uuid import uuid4
-
-import pandas as pd
+from functools import wraps
 
-from validmind.ai.test_descriptions import get_description_metadata
-from validmind.errors import MissingRequiredTestInputError
 from validmind.logging import get_logger
-from validmind.vm_models import (
-    Metric,
-    MetricResult,
-    ResultSummary,
-    ResultTable,
-    ResultTableMetadata,
-    VMDataset,
-    VMModel,
-)
-from validmind.vm_models.figure import (
-    Figure,
-    is_matplotlib_figure,
-    is_plotly_figure,
-    is_png_image,
-)
-from validmind.vm_models.test.result_wrapper import MetricResultWrapper
 
 from ._store import test_store
+from .load import load_test
 
 logger = get_logger(__name__)
 
 
-_input_type_map = {
-    "dataset": VMDataset,
-    "datasets": List[VMDataset],
-    "model": VMModel,
-    "models": List[VMModel],
-}
-
-
-def _inspect_signature(test_func: callable):
-    inputs = {}
-    params = {}
-
-    for name, arg in inspect.signature(test_func).parameters.items():
-        if name in _input_type_map:
-            inputs[name] = {
-                "type": _input_type_map[name],
-            }
-        else:
-            params[name] = {
-                "type": arg.annotation,
-                "default": (
-                    arg.default if arg.default is not inspect.Parameter.empty else None
-                ),
-            }
-
-    return inputs, params
-
-
-def _build_result(  # noqa: C901
-    results: Union[Any, Tuple[Any, ...]],
-    test_id: str,
-    inputs: List[str],
-    params: Dict[str, Any],
-    description: str = None,
-    output_template: str = None,
-    generate_description: bool = True,
-):
-    ref_id = str(uuid4())
-    figure_metadata = {
-        "_type": "metric",
-        "_name": test_id,
-        "_ref_id": ref_id,
-    }
-
-    tables = []
-    figures = []
-    scalars = []
-
-    def process_result_item(item):
-        # TOOD: build out a more robust/extensible system for this
-        # TODO: custom type handlers would be really cool
-
-        # unit metrics (scalar values) - for now only one per test
-        if isinstance(item, int) or isinstance(item, float):
-            if scalars:
-                raise ValueError("Only one unit metric may be returned per test.")
-            scalars.append(item)
-
-        # plots
-        elif isinstance(item, Figure):
-            figures.append(item)
-        elif is_matplotlib_figure(item) or is_plotly_figure(item) or is_png_image(item):
-            figures.append(
-                Figure(
-                    key=f"{test_id}:{len(figures) + 1}",
-                    figure=item,
-                    metadata=figure_metadata,
-                )
-            )
-
-        # tables
-        elif isinstance(item, list) or isinstance(item, pd.DataFrame):
-            tables.append(ResultTable(data=item))
-        elif isinstance(item, dict):
-            for table_name, table in item.items():
-                if not isinstance(table, list) and not isinstance(table, pd.DataFrame):
-                    raise ValueError(
-                        f"Invalid table format: {table_name} must be a list or DataFrame"
-                    )
-
-                tables.append(
-                    ResultTable(
-                        data=table,
-                        metadata=ResultTableMetadata(title=table_name),
-                    )
-                )
-
-        else:
-            raise ValueError(f"Invalid return type: {type(item)}")
-
-    # if the results are a tuple, process each item as a separate result
-    if isinstance(results, tuple):
-        for item in results:
-            process_result_item(item)
-    else:
-        process_result_item(results)
-
-    metric_inputs = [
-        sub_i.input_id if hasattr(sub_i, "input_id") else sub_i
-        for i in inputs
-        for sub_i in (i if isinstance(i, list) else [i])
-    ]
-
-    return MetricResultWrapper(
-        result_id=test_id,
-        scalar=scalars[0] if scalars else None,
-        metric=(
-            MetricResult(
-                key=test_id,
-                ref_id=ref_id,
-                value="Empty",
-                summary=ResultSummary(results=tables),
-            )
-            if tables or figures  # if tables or figures than its a traditional metric
-            else None
-        ),
-        figures=figures,
-        result_metadata=(
-            [
-                get_description_metadata(
-                    test_id=test_id,
-                    default_description=description,
-                    summary=ResultSummary(results=tables).serialize(),
-                    figures=figures,
-                    should_generate=generate_description,
-                )
-            ]
-            if tables or figures
-            else None
-        ),
-        inputs=metric_inputs,
-        params=params,
-        output_template=output_template,
-    )
-
-
-def _get_run_method(func, func_inputs, func_params):
-    def run(self: Metric):
-        input_kwargs = {}  # map function inputs (`dataset` etc) to actual objects
-        input_ids = []  # store input_ids used so they can be logged
-        for key in func_inputs.keys():
-            try:
-                input_kwargs[key] = getattr(self.inputs, key)
-                if isinstance(input_kwargs[key], list):
-                    input_ids.extend([i.input_id for i in input_kwargs[key]])
-                else:
-                    input_ids.append(input_kwargs[key].input_id)
-            except AttributeError:
-                raise MissingRequiredTestInputError(f"Missing required input: {key}.")
-
-        param_kwargs = {
-            key: self.params.get(key, func_params[key]["default"])
-            for key in func_params.keys()
-        }
-
-        raw_results = func(**input_kwargs, **param_kwargs)
-
-        self.result = _build_result(
-            results=raw_results,
-            test_id=self.test_id,
-            description=inspect.getdoc(self),
-            inputs=input_ids,
-            params=param_kwargs,
-            output_template=self.output_template,
-            generate_description=self.generate_description,
-        )
-
-        return self.result
-
-    return run
+def _get_save_func(func, test_id):
+    """Helper function to save a decorated function to a file
 
+    Useful when a custom test function has been created inline in a notebook or
+    interactive session and needs to be saved to a file so it can be added to a
+    test library.
+    """
 
-def _get_save_func(func, test_id):
     def save(root_folder=".", imports=None):
         parts = test_id.split(".")
 
@@ -270,34 +82,26 @@ def _get_save_func(func, test_id):
     return save
 
 
-def metric(func_or_id):
-    """
-    DEPRECATED, use @vm.test instead
-    """
-    # print a deprecation notice and call the test() function instead
-    logger.warning(
-        "The @vm.metric decorator is deprecated and will be removed in a future release. "
-        "Please use @vm.test instead."
-    )
-    return test(func_or_id)
+def test(func_or_id):
+    """Decorator for creating and registering custom tests
 
+    This decorator registers the function it wraps as a test function within ValidMind
+    under the provided ID. Once decorated, the function can be run using the
+    `validmind.tests.run_test` function.
 
-def test(func_or_id):
-    """Decorator for creating and registering metrics with the ValidMind framework.
+    The function can take two different types of arguments:
 
-    Creates a metric object and registers it with ValidMind under the provided ID. If
-    no ID is provided, the function name will be used as to build one. So if the
-    function name is `my_metric`, the metric will be registered under the ID
-    `validmind.custom_metrics.my_metric`.
+    - Inputs: ValidMind model or dataset (or list of models/datasets). These arguments
+      must use the following names: `model`, `models`, `dataset`, `datasets`.
+    - Parameters: Any additional keyword arguments of any type (must have a default
+      value) that can have any name.
 
-    This decorator works by creating a new `Metric` class will be created whose `run`
-    method calls the decorated function. This function should take as arguments the
-    inputs it requires (`dataset`, `datasets`, `model`, `models`) followed by any
-    parameters. It can return any number of the following types:
+    The function should return one of the following types:
 
     - Table: Either a list of dictionaries or a pandas DataFrame
     - Plot: Either a matplotlib figure or a plotly figure
-    - Scalar: A single number or string
+    - Scalar: A single number (int or float)
+    - Boolean: A single boolean value indicating whether the test passed or failed
 
     The function may also include a docstring. This docstring will be used and logged
     as the metric's description.
@@ -312,30 +116,17 @@ def test(func_or_id):
 
     def decorator(func):
         test_id = func_or_id or f"validmind.custom_metrics.{func.__name__}"
+        test_func = load_test(test_id, func, reload=True)
+        test_store.register_test(test_id, test_func)
 
-        inputs, params = _inspect_signature(func)
-        description = inspect.getdoc(func)
-        tasks = getattr(func, "__tasks__", [])
-        tags = getattr(func, "__tags__", [])
-
-        metric_class = type(
-            func.__name__,
-            (Metric,),
-            {
-                "run": _get_run_method(func, inputs, params),
-                "required_inputs": list(inputs.keys()),
-                "default_params": {k: v["default"] for k, v in params.items()},
-                "__doc__": description,
-                "tasks": tasks,
-                "tags": tags,
-            },
-        )
-        test_store.register_custom_test(test_id, metric_class)
+        @wraps(test_func)
+        def wrapper(*args, **kwargs):
+            return test_func(*args, **kwargs)
 
         # special function to allow the function to be saved to a file
-        func.save = _get_save_func(func, test_id)
+        wrapper.save = _get_save_func(test_func, test_id)
 
-        return func
+        return wrapper
 
     if callable(func_or_id):
         return decorator(func_or_id)
@@ -344,10 +135,10 @@ def test(func_or_id):
 
 
 def tasks(*tasks):
-    """Decorator for specifying the task types that a metric is designed for.
+    """Decorator for specifying the task types that a test is designed for.
 
     Args:
-        *tasks: The task types that the metric is designed for.
+        *tasks: The task types that the test is designed for.
     """
 
     def decorator(func):
@@ -358,10 +149,10 @@ def tasks(*tasks):
 
 
 def tags(*tags):
-    """Decorator for specifying tags for a metric.
+    """Decorator for specifying tags for a test.
 
     Args:
-        *tags: The tags to apply to the metric.
+        *tags: The tags to apply to the test.
     """
 
     def decorator(func):