Source code for fleche.metadata

from abc import ABC, abstractmethod
from dataclasses import dataclass
import time
from typing import Any, TypeAlias

from .call import Call

# Values produced by MetaData.pre/post must be JSON-serializable.
# This alias documents the expected shape and helps static type checkers.


JSONValue: TypeAlias = str | int | float | bool | None | list["JSONValue"] | dict[str, "JSONValue"]





class MetaData(ABC):
    """Abstract base class for defining metadata types.

    Implementations must return only JSON-serializable values from pre() and post().
    That means scalars (str, int, float, bool, None), lists of those, or dicts with str keys
    and JSON-serializable values.
    """


    def pre(self, call: Call) -> dict[str, JSONValue]:
        """
        Hook for collecting metadata before the function execution.

        Args:
            call (Call): The call object of the decorated function.

        Returns:
            dict[str, JSONValue]: A flat dictionary of JSON-serializable metadata collected before execution.
        """
        return {}




    def post(self, pre: dict[str, JSONValue], call: Call) -> dict[str, JSONValue]:
        """
        Hook for collecting metadata after the function execution.

        The return value of the function is available on the `call.result` attribute.

        Args:
            pre (dict[str, JSONValue]): Metadata collected during the pre-execution phase.
            call (Call): The call object of the decorated function.

        Returns:
            dict[str, JSONValue]: A flat dictionary of JSON-serializable metadata collected after execution.
        """
        return {}


    @property
    @abstractmethod


    def keys(self) -> dict[str, type]:
        """
        Defines the schema of the metadata, mapping metadata keys to their expected types.

        Returns:
            dict[str, type]: A dictionary representing the metadata schema.
        """
        ...


    @property
    @abstractmethod


    def name(self) -> str:
        """
        The unique name of this metadata type.

        Returns:
            str: The name of the metadata type.
        """
        ...






class Runtime(MetaData):
    """Metadata type for capturing runtime information.

    Keys:
        timestart (float): The timestamp when the execution started.
        timestop (float): The timestamp when the execution stopped.
        walltime (float): The total execution time in seconds.

    Notes:
        Values are JSON-serializable.
    """


    def pre(self, call: Call) -> dict[str, Any]:
        """
        Records the start time before function execution.
        """
        return {'timestart': time.time()}




    def post(self, pre: dict[str, Any], call: Call) -> dict[str, Any]:
        """
        Records the stop time and calculates the wall time after function execution.
        """
        return {
            'timestop': (t := time.time()),
            'walltime': t - pre['timestart'],
        }




    name: str = 'runtime'



    keys: dict[str, type] = {
            'timestart': float,
            'timestop': float,
            'walltime': float,
    }




@dataclass


class Tags(MetaData):
    """Metadata type for storing arbitrary tags.

    For each key in the ``tags`` dictionary, a new metadata column is created.

    Keys:
        tags (dict): A dictionary of user-defined tags.

    Notes:
        Tag values must be JSON-serializable.
    """


    tags: dict[str, Any]




    def pre(self, call: Call) -> dict[str, Any]:
        return self.tags.copy()




    name: str = "tags"


    @property


    def keys(self):
        return {k: type(v) for k, v in self.tags.items()}