venv/lib/python3.11/site-packages/pandas/tests/frame/test_iteration.py

import datetime

import numpy as np
import pytest

from pandas.compat import (
    IS64,
    is_platform_windows,
)

from pandas import (
    Categorical,
    DataFrame,
    Series,
    date_range,
)
import pandas._testing as tm


class TestIteration:
    def test_keys(self, float_frame):
        assert float_frame.keys() is float_frame.columns

    def test_iteritems(self):
        df = DataFrame([[1, 2, 3], [4, 5, 6]], columns=["a", "a", "b"])
        for k, v in df.items():
            assert isinstance(v, DataFrame._constructor_sliced)

    def test_items(self):
        # GH#17213, GH#13918
        cols = ["a", "b", "c"]
        df = DataFrame([[1, 2, 3], [4, 5, 6]], columns=cols)
        for c, (k, v) in zip(cols, df.items()):
            assert c == k
            assert isinstance(v, Series)
            assert (df[k] == v).all()

    def test_items_names(self, float_string_frame):
        for k, v in float_string_frame.items():
            assert v.name == k

    def test_iter(self, float_frame):
        assert list(float_frame) == list(float_frame.columns)

    def test_iterrows(self, float_frame, float_string_frame):
        for k, v in float_frame.iterrows():
            exp = float_frame.loc[k]
            tm.assert_series_equal(v, exp)

        for k, v in float_string_frame.iterrows():
            exp = float_string_frame.loc[k]
            tm.assert_series_equal(v, exp)

    def test_iterrows_iso8601(self):
        # GH#19671
        s = DataFrame(
            {
                "non_iso8601": ["M1701", "M1802", "M1903", "M2004"],
                "iso8601": date_range("2000-01-01", periods=4, freq="ME"),
            }
        )
        for k, v in s.iterrows():
            exp = s.loc[k]
            tm.assert_series_equal(v, exp)

    def test_iterrows_corner(self):
        # GH#12222
        df = DataFrame(
            {
                "a": [datetime.datetime(2015, 1, 1)],
                "b": [None],
                "c": [None],
                "d": [""],
                "e": [[]],
                "f": [set()],
                "g": [{}],
            }
        )
        expected = Series(
            [datetime.datetime(2015, 1, 1), None, None, "", [], set(), {}],
            index=list("abcdefg"),
            name=0,
            dtype="object",
        )
        _, result = next(df.iterrows())
        tm.assert_series_equal(result, expected)

    def test_itertuples(self, float_frame):
        for i, tup in enumerate(float_frame.itertuples()):
            ser = DataFrame._constructor_sliced(tup[1:])
            ser.name = tup[0]
            expected = float_frame.iloc[i, :].reset_index(drop=True)
            tm.assert_series_equal(ser, expected)

    def test_itertuples_index_false(self):
        df = DataFrame(
            {"floats": np.random.default_rng(2).standard_normal(5), "ints": range(5)},
            columns=["floats", "ints"],
        )

        for tup in df.itertuples(index=False):
            assert isinstance(tup[1], int)

    def test_itertuples_duplicate_cols(self):
        df = DataFrame(data={"a": [1, 2, 3], "b": [4, 5, 6]})
        dfaa = df[["a", "a"]]

        assert list(dfaa.itertuples()) == [(0, 1, 1), (1, 2, 2), (2, 3, 3)]

        # repr with int on 32-bit/windows
        if not (is_platform_windows() or not IS64):
            assert (
                repr(list(df.itertuples(name=None)))
                == "[(0, 1, 4), (1, 2, 5), (2, 3, 6)]"
            )

    def test_itertuples_tuple_name(self):
        df = DataFrame(data={"a": [1, 2, 3], "b": [4, 5, 6]})
        tup = next(df.itertuples(name="TestName"))
        assert tup._fields == ("Index", "a", "b")
        assert (tup.Index, tup.a, tup.b) == tup
        assert type(tup).__name__ == "TestName"

    def test_itertuples_disallowed_col_labels(self):
        df = DataFrame(data={"def": [1, 2, 3], "return": [4, 5, 6]})
        tup2 = next(df.itertuples(name="TestName"))
        assert tup2 == (0, 1, 4)
        assert tup2._fields == ("Index", "_1", "_2")

    @pytest.mark.parametrize("limit", [254, 255, 1024])
    @pytest.mark.parametrize("index", [True, False])
    def test_itertuples_py2_3_field_limit_namedtuple(self, limit, index):
        # GH#28282
        df = DataFrame([{f"foo_{i}": f"bar_{i}" for i in range(limit)}])
        result = next(df.itertuples(index=index))
        assert isinstance(result, tuple)
        assert hasattr(result, "_fields")

    def test_sequence_like_with_categorical(self):
        # GH#7839
        # make sure can iterate
        df = DataFrame(
            {"id": [1, 2, 3, 4, 5, 6], "raw_grade": ["a", "b", "b", "a", "a", "e"]}
        )
        df["grade"] = Categorical(df["raw_grade"])

        # basic sequencing testing
        result = list(df.grade.values)
        expected = np.array(df.grade.values).tolist()
        tm.assert_almost_equal(result, expected)

        # iteration
        for t in df.itertuples(index=False):
            str(t)

        for row, s in df.iterrows():
            str(s)

        for c, col in df.items():
            str(col)