Namespace registration

import random

import numpy as np
import pandas as pd
import polars as pl
import sklearn.datasets

df = pl.from_pandas(sklearn.datasets.load_digits(as_frame=True)["frame"].reset_index())

@pl.api.register_dataframe_namespace("random")
class CustomDataFrameNamespace:
    def __init__(self, df: pl.DataFrame):
        self._df = df
    
    def sample_consecutive(self, n: int, seed: int = None) -> pl.DataFrame:
        """
        Sample n consecutive rows from the DataFrame.
        
        Parameters:
        -----------
        n : int
            Number of consecutive rows to sample
        seed : int, optional
            Random seed for reproducibility
            
        Returns:
        --------
        pl.DataFrame
            DataFrame containing n consecutive rows
        """
        if n <= 0:
            raise ValueError("n must be positive")

        df_len = self._df.height
        if n > df_len:
            raise ValueError(f"n ({n}) cannot be larger than DataFrame height ({df_len})")       

        if seed is not None:
            random.seed(seed)

        # Calculate the maximum starting index
        max_start_idx = df_len - n
        # Randomly choose starting index
        start_idx = random.randint(0, max_start_idx)
        
        # Return consecutive rows using slice
        return self._df.slice(start_idx, n)

df.random.sample_consecutive(5)

# Using LazyFrame namespace
@pl.api.register_lazyframe_namespace("random")
class CustomLazyFrameNamespace:
    def __init__(self, lf: pl.LazyFrame):
        self._lf = lf
    
    def sample_consecutive(self, n: int, seed: int = None) -> pl.LazyFrame:
        """Sample n consecutive rows from the LazyFrame."""
        if n <= 0:
            raise ValueError("n must be positive")        
     
        # For LazyFrame, we need to collect first to get the length
        # In practice, you might want to handle this differently
        df_len = self._lf.collect().height
        if n > df_len:
            raise ValueError(f"n ({n}) cannot be larger than DataFrame height ({df_len})")

        if seed is not None:
            random.seed(seed)

        max_start_idx = df_len - n
        start_idx = random.randint(0, max_start_idx)
        
        return self._lf.slice(start_idx, n)

@pl.api.register_series_namespace("math")
class MathShortcuts:
    def __init__(self, s: pl.Series) -> None:
        self._s = s

    def square(self) -> pl.Series:
        return self._s * self._s

    def cube(self) -> pl.Series:
        return self._s * self._s * self._s

df["index"].math.square()

@pd.api.extensions.register_dataframe_accessor("random")
class CustomAccessor:
    def __init__(self, pandas_obj):
        self._obj = pandas_obj
    
    def sample_consecutive(self, n, random_state=None):
        """
        Sample n consecutive rows from the DataFrame.
        
        Parameters:
        -----------
        n : int
            Number of consecutive rows to sample
        random_state : int, optional
            Random seed for reproducibility
            
        Returns:
        --------
        pandas.DataFrame
            DataFrame containing n consecutive rows
        """
        if n <= 0:
            raise ValueError("n must be positive")
        
        if n > len(self._obj):
            raise ValueError(f"n ({n}) cannot be larger than DataFrame length ({len(self._obj)})")
        
        # Set random seed if provided
        if random_state is not None:
            np.random.seed(random_state)
        
        # Calculate the maximum starting index
        max_start_idx = len(self._obj) - n
        
        # Randomly choose starting index
        start_idx = np.random.randint(0, max_start_idx + 1)
        
        # Return consecutive rows
        return self._obj.iloc[start_idx:start_idx + n]

df = pd.DataFrame({
    'A': range(1, 21),
    'B': [f'item_{i}' for i in range(1, 21)],
    'C': np.random.randn(20)
})

df.head()

df.random.sample_consecutive(5, random_state=42)

df.random.sample_consecutive(3, random_state=123)