from typing import Callable
import functools
from typing import Any
import time


def describe(func: Callable) -> Callable:
    """Describe the shape of the input shape, output shape, and time of a pandas pipe function."""

    @functools.wraps(func)
    def wrapper(*args: Any, **kwargs: Any) -> Any:
        pre_shape = str(args[0].shape)
        n_rows = args[0].shape[0]
        start_ts = time.time()

        result = func(*args, **kwargs)

        try:
            post_shape = str(result.shape)
        except AttributeError:
            post_shape = 'NA'

        end_ts = time.time()
        sec = end_ts - start_ts
        sec_per_mille = (sec / n_rows) * 1000

        print(f'{func.__name__}')
        print(f'  Shape: {pre_shape} → {post_shape}')
        print(f'  Time: {sec:.2f}s ({sec_per_mille:.2f}s / 1k)')
        return result

    return wrapper