test_pandas.py

# Copyright 2023 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import inspect
import re
import sys
import unittest.mock as mock

import pandas as pd
import pytest

import bigframes.core.global_session
import bigframes.dataframe
import bigframes.pandas as bpd
import bigframes.session

leading_whitespace = re.compile(r"^\s+", flags=re.MULTILINE)


def all_session_methods():
    session_attributes = set(
        attribute
        for attribute in dir(bigframes.session.Session)
        if not attribute.startswith("_")
    )
    session_attributes.remove("close")
    # streaming isn't in pandas
    session_attributes.remove("read_gbq_table_streaming")

    for attribute in sorted(session_attributes):
        session_method = getattr(bigframes.session.Session, attribute)
        if not callable(session_method):
            continue  # pragma: no cover
        yield attribute


@pytest.mark.parametrize(
    [
        "method_name",
    ],
    [(method_name,) for method_name in all_session_methods()],
)
def test_method_matches_session(method_name: str):
    if sys.version_info < (3, 10):
        pytest.skip(
            "Need Python 3.10 to reconcile deferred annotations."
        )  # pragma: no cover

    session_method = getattr(bigframes.session.Session, method_name)
    session_doc = inspect.getdoc(session_method)
    assert session_doc is not None, "docstrings are required"

    pandas_method = getattr(bigframes.pandas, method_name)
    pandas_doc = inspect.getdoc(pandas_method)
    assert pandas_doc is not None, "docstrings are required"
    assert re.sub(leading_whitespace, "", pandas_doc) == re.sub(
        leading_whitespace, "", session_doc
    )

    # Add `eval_str = True` so that deferred annotations are turned into their
    # corresponding type objects. Need Python 3.10 for eval_str parameter.
    session_signature = inspect.signature(
        session_method,
        eval_str=True,
        globals={**vars(bigframes.session), **{"dataframe": bigframes.dataframe}},
    )
    pandas_signature = inspect.signature(pandas_method, eval_str=True)
    assert [
        # Kind includes position, which will be an offset.
        parameter.replace(kind=inspect.Parameter.POSITIONAL_ONLY)
        for parameter in pandas_signature.parameters.values()
    ] == [
        # Kind includes position, which will be an offset.
        parameter.replace(kind=inspect.Parameter.POSITIONAL_ONLY)
        for parameter in session_signature.parameters.values()
        # Don't include the first parameter, which is `self: Session`
    ][
        1:
    ]
    assert pandas_signature.return_annotation == session_signature.return_annotation


def test_cut_raises_with_labels():
    with pytest.raises(
        NotImplementedError,
        match="The 'labels' parameter must be either False or None.",
    ):
        mock_series = mock.create_autospec(bigframes.pandas.Series, instance=True)
        bigframes.pandas.cut(mock_series, 4, labels=["a", "b", "c", "d"])


@pytest.mark.parametrize(
    ("bins",),
    (
        (0,),
        (-1,),
    ),
)
def test_cut_raises_with_invalid_bins(bins: int):
    with pytest.raises(ValueError, match="`bins` should be a positive integer."):
        mock_series = mock.create_autospec(bigframes.pandas.Series, instance=True)
        bigframes.pandas.cut(mock_series, bins, labels=False)


def test_pandas_attribute():
    assert bpd.NA is pd.NA
    assert bpd.BooleanDtype is pd.BooleanDtype
    assert bpd.Float64Dtype is pd.Float64Dtype
    assert bpd.Int64Dtype is pd.Int64Dtype
    assert bpd.StringDtype is pd.StringDtype
    assert bpd.ArrowDtype is pd.ArrowDtype