function decorator to extend polars with new method via dataframe namespace accessors

new_method.py

"""Example of adding new methods to Polars DataFrames using decorators and namespace accessors.

Examples
--------
Make a new method that doubles all values in a DataFrame:

.. code-block:: python

    import polars as pl


    @new_method
    def double(df: pl.DataFrame) -> pl.DataFrame:
        return df * 2


    df = pl.DataFrame(
        {
            "a": [1, 2, 3],
            "b": [4, 5, 6],
        }
    )

    df.double()

References:

- https://docs.pola.rs/api/python/stable/reference/api.html

"""

import polars as pl
from functools import wraps


def new_method(func):
    """A decorator to add a new method to DataFrames."""

    name = func.__name__

    @pl.api.register_dataframe_namespace(name)
    class NewNameSpace:
        def __init__(self, df: pl.DataFrame):
            self._df = df

        @wraps(func)
        def __call__(self, *args, **kwargs):
            return func(self._df, *args, **kwargs)

    return func