How to use siuba - 10 common examples

def f(x):
        return GroupByAgg.from_result(x.mean(), x)

def _f_grouped(x) -> GroupByAgg:
        return GroupByAgg.from_result(x.mean() + 10, x)

def f(x) -> Any:
        return GroupByAgg.from_result(x.mean() + 10, x)

def test_preserves_original_keys():
    df = data_frame(x = [1,2], y = 2, z = [1,2])
    pv = pivot_longer(df, _["y":"z"])

    assert pv.columns.tolist() == ["x", "name", "value"]
    assert assert_series_equal(
            pv["x"],
            pd.Series(df["x"].repeat(2))
            )

from siuba.tests.helpers import data_frame
import pandas as pd

from siuba.experimental.pd_groups.translate import method_agg_op, method_el_op, method_el_op2
#TODO: 
#  - what if they have mandatory, non-data args?
#  - support accessor methods like _.x.str.upper()
#  - support .expanding and .rolling


data_dt = data_frame(
    g = ['a', 'a', 'b', 'b'],
    x = pd.to_datetime(["2019-01-01 01:01:01", "2020-04-08 02:02:02", "2021-07-15 03:03:03", "2022-10-22 04:04:04"])
    )

data_str = data_frame(
    g = ['a', 'a', 'b', 'b'],
    x = ['abc', 'cde', 'fg', 'h']
    )

data_default = data_frame(
    g = ['a', 'a', 'b', 'b'],
    x = [10, 11, 12, 13],
    y = [1,2,3,4]
    )

data = {
    'dt': data_dt,
    'str': data_str,
    None: data_default
}

def test_spec_add_multi_columns():
    df = data_frame(x = [1,2], y = [3,4])

    # TODO: is this the right format for a spec
    sp = data_frame(_name = ["x", "y"], _value = "v", a = 1, b = 2)
    pv = pivot_longer_spec(df, spec = sp)

    assert pv.columns.tolist() == ["a", "b", "v"]

def test_spec_add_multi_columns():
    df = data_frame(x = [1,2], y = [3,4])

    # TODO: is this the right format for a spec
    sp = data_frame(_name = ["x", "y"], _value = "v", a = 1, b = 2)
    pv = pivot_longer_spec(df, spec = sp)

    assert pv.columns.tolist() == ["a", "b", "v"]

#  - what if they have mandatory, non-data args?
#  - support accessor methods like _.x.str.upper()
#  - support .expanding and .rolling


data_dt = data_frame(
    g = ['a', 'a', 'b', 'b'],
    x = pd.to_datetime(["2019-01-01 01:01:01", "2020-04-08 02:02:02", "2021-07-15 03:03:03", "2022-10-22 04:04:04"])
    )

data_str = data_frame(
    g = ['a', 'a', 'b', 'b'],
    x = ['abc', 'cde', 'fg', 'h']
    )

data_default = data_frame(
    g = ['a', 'a', 'b', 'b'],
    x = [10, 11, 12, 13],
    y = [1,2,3,4]
    )

data = {
    'dt': data_dt,
    'str': data_str,
    None: data_default
}


# Test translator =============================================================

from pandas.testing import assert_frame_equal, assert_series_equal
from siuba.experimental.pd_groups.groupby import GroupByAgg, SeriesGroupBy

def test_pivot_all_cols_to_long():
    "can pivot all cols to long"

    src = data_frame(x = [1,2], y = [3,4])
    dst = data_frame(name = ["x", "y", "x", "y"], value = [1, 3, 2, 4])
    
    res = pivot_longer(src, _["x":"y"])

    assert_frame_sort_equal(res, dst)

def test_pivot_all_cols_to_long():
    "can pivot all cols to long"

    src = data_frame(x = [1,2], y = [3,4])
    dst = data_frame(name = ["x", "y", "x", "y"], value = [1, 3, 2, 4])
    
    res = pivot_longer(src, _["x":"y"])

    assert_frame_sort_equal(res, dst)