Expand numeric range in pandas

expand_numeric_range.py

from typing import Union
import pandas as pd

def expand_numeric_range(
    frame: pd.DataFrame,
    colname: str,
    start: str,
    stop: str,
    step: Union[str, int] = 1,
) -> pd.DataFrame:
    """Expands a numeric range along an axis of a DataFrame
    
    Parameters
    ----------
    frame: DataFrame
        The DataFrame to which this transformation applies.
    colname: str
        The name of the new column that is created.
    start: str
        The name of the column that contains the start of the range.
    stop: str
        The name of the column that contains the end of the range (inclusive).
    step: str | int
        The amount of incrementation, or the name of the column containing it.
        
    Returns
    -------
    Dataframe
        The transformed dataframe
        
    Examples:
    
    >>> df = pd.DataFrame([(1, 'Alice', 2009, 2010),
    ...                    (2, 'Bob', 2010, 2012),
    ...                    (3, 'Carol', 2008, 2010)],
    ...                   columns=('id', 'name', 'year_start' , 'year_end'))
    >>> df.pipe(
    ...     expand_numeric_range,
    ...     colname='year',
    ...     start='year_start',
    ...     stop='year_end',
    ... )
       id   name  year
    0   1  Alice  2009
    0   1  Alice  2010
    1   2    Bob  2010
    1   2    Bob  2011
    1   2    Bob  2012
    2   3  Carol  2008
    2   3  Carol  2009
    2   3  Carol  2010
    """
    
    new_frame = (frame
        .assign(
            __step=(lambda df: df[step]) if isinstance(step, str) else step,
            __new_value=lambda df: df.apply(
                lambda s: range(s[start], s[stop]+1, s['__step']),
                axis='columns',
            )
        )
        .drop([start, stop, '__step'], axis='columns')
        .rename({'__new_value': colname}, axis='columns')
        .explode(colname)
    )
    
    return new_frame