Repository: pandas-dev/pandas
File: pandas/tests/test_col.py (Python)

Code before:


@pytest.mark.parametrize(
    ("expr", "expected_values", "expected_str"),
    [
        (pd.col("a"), [1, 2], "col('a')"),
        (pd.col("a") * 2, [2, 4], "(col('a') * 2)"),
        (pd.col("a").sum(), [3, 3], "col('a').sum()"),
        (pd.col("a") + 1, [2, 3], "(col('a') + 1)"),
        (1 + pd.col("a"), [2, 3], "(1 + col('a'))"),
        (pd.col("a") - 1, [0, 1], "(col('a') - 1)"),
        (1 - pd.col("a"), [0, -1], "(1 - col('a'))"),
        (pd.col("a") * 1, [1, 2], "(col('a') * 1)"),
        (1 * pd.col("a"), [1, 2], "(1 * col('a'))"),
        (pd.col("a") / 1, [1.0, 2.0], "(col('a') / 1)"),
        (1 / pd.col("a"), [1.0, 0.5], "(1 / col('a'))"),
        (pd.col("a") // 1, [1, 2], "(col('a') // 1)"),
        (1 // pd.col("a"), [1, 0], "(1 // col('a'))"),
        (pd.col("a") % 1, [0, 0], "(col('a') % 1)"),
        (1 % pd.col("a"), [0, 1], "(1 % col('a'))"),
        (pd.col("a") > 1, [False, True], "(col('a') > 1)"),
        (pd.col("a") >= 1, [True, True], "(col('a') >= 1)"),
        (pd.col("a") < 1, [False, False], "(col('a') < 1)"),
        (pd.col("a") <= 1, [True, False], "(col('a') <= 1)"),
        (pd.col("a") == 1, [True, False], "(col('a') == 1)"),
        (np.log(pd.col("a")), [0.0, 0.6931471805599453], "log(col('a'))"),
        (np.divide(pd.col("a"), pd.col("a")), [1.0, 1.0], "divide(col('a'), col('a'))"),
    ],
)
def test_col_simple(
    expr: Expression, expected_values: list[object], expected_str: str
) -> None:
    df = pd.DataFrame({"a": [1, 2], "b": [3, 4]})
    result = df.assign(c=expr)
    expected = pd.DataFrame({"a": [1, 2], "b": [3, 4], "c": expected_values})
    tm.assert_frame_equal(result, expected)
    assert str(expr) == expected_str


@pytest.mark.parametrize(
    ("expr", "expected_values", "expected_str"),
    [
        (pd.col("a").dt.year, [2020], "col('a').dt.year"),
        (pd.col("a").dt.strftime("%B"), ["January"], "col('a').dt.strftime('%B')"),
        (pd.col("b").str.upper(), ["FOO"], "col('b').str.upper()"),
    ],
)
def test_namespaces(
    expr: Expression, expected_values: list[object], expected_str: str
) -> None:
    df = pd.DataFrame({"a": [datetime(2020, 1, 1)], "b": ["foo"]})

Review comment (suggestion):
`log` is dangerous to use in a test because the floating point value could be different on different platforms.  

Maybe use `np.min()` instead.


Code after:


@pytest.mark.parametrize(
    ("expr", "expected_values", "expected_str"),
    [
        (pd.col("a"), [1, 2], "col('a')"),
        (pd.col("a") * 2, [2, 4], "(col('a') * 2)"),
        (pd.col("a").sum(), [3, 3], "col('a').sum()"),
        (pd.col("a") + 1, [2, 3], "(col('a') + 1)"),
        (1 + pd.col("a"), [2, 3], "(1 + col('a'))"),
        (pd.col("a") - 1, [0, 1], "(col('a') - 1)"),
        (1 - pd.col("a"), [0, -1], "(1 - col('a'))"),
        (pd.col("a") * 1, [1, 2], "(col('a') * 1)"),
        (1 * pd.col("a"), [1, 2], "(1 * col('a'))"),
        (pd.col("a") / 1, [1.0, 2.0], "(col('a') / 1)"),
        (1 / pd.col("a"), [1.0, 0.5], "(1 / col('a'))"),
        (pd.col("a") // 1, [1, 2], "(col('a') // 1)"),
        (1 // pd.col("a"), [1, 0], "(1 // col('a'))"),
        (pd.col("a") % 1, [0, 0], "(col('a') % 1)"),
        (1 % pd.col("a"), [0, 1], "(1 % col('a'))"),
        (pd.col("a") > 1, [False, True], "(col('a') > 1)"),
        (pd.col("a") >= 1, [True, True], "(col('a') >= 1)"),
        (pd.col("a") < 1, [False, False], "(col('a') < 1)"),
        (pd.col("a") <= 1, [True, False], "(col('a') <= 1)"),
        (pd.col("a") == 1, [True, False], "(col('a') == 1)"),
        (np.power(pd.col("a"), 2), [1, 4], "power(col('a'), 2)"),
        (np.divide(pd.col("a"), pd.col("a")), [1.0, 1.0], "divide(col('a'), col('a'))"),
    ],
)
def test_col_simple(
    expr: Expression, expected_values: list[object], expected_str: str
) -> None:
    df = pd.DataFrame({"a": [1, 2], "b": [3, 4]})
    result = df.assign(c=expr)
    expected = pd.DataFrame({"a": [1, 2], "b": [3, 4], "c": expected_values})
    tm.assert_frame_equal(result, expected)
    assert str(expr) == expected_str


@pytest.mark.parametrize(
    ("expr", "expected_values", "expected_str"),
    [
        (pd.col("a").dt.year, [2020], "col('a').dt.year"),
        (pd.col("a").dt.strftime("%B"), ["January"], "col('a').dt.strftime('%B')"),
        (pd.col("b").str.upper(), ["FOO"], "col('b').str.upper()"),
    ],
)
def test_namespaces(
    expr: Expression, expected_values: list[object], expected_str: str
) -> None:
    df = pd.DataFrame({"a": [datetime(2020, 1, 1)], "b": ["foo"]})