venv/lib/python3.11/site-packages/pandas/tests/io/test_stata.py

import bz2
import datetime as dt
from datetime import datetime
import gzip
import io
import os
import struct
import tarfile
import zipfile

import numpy as np
import pytest

import pandas.util._test_decorators as td

import pandas as pd
from pandas import CategoricalDtype
import pandas._testing as tm
from pandas.core.frame import (
    DataFrame,
    Series,
)

from pandas.io.parsers import read_csv
from pandas.io.stata import (
    CategoricalConversionWarning,
    InvalidColumnName,
    PossiblePrecisionLoss,
    StataMissingValue,
    StataReader,
    StataWriter,
    StataWriterUTF8,
    ValueLabelTypeMismatch,
    read_stata,
)


@pytest.fixture
def mixed_frame():
    return DataFrame(
        {
            "a": [1, 2, 3, 4],
            "b": [1.0, 3.0, 27.0, 81.0],
            "c": ["Atlanta", "Birmingham", "Cincinnati", "Detroit"],
        }
    )


@pytest.fixture
def parsed_114(datapath):
    dta14_114 = datapath("io", "data", "stata", "stata5_114.dta")
    parsed_114 = read_stata(dta14_114, convert_dates=True)
    parsed_114.index.name = "index"
    return parsed_114


class TestStata:
    def read_dta(self, file):
        # Legacy default reader configuration
        return read_stata(file, convert_dates=True)

    def read_csv(self, file):
        return read_csv(file, parse_dates=True)

    @pytest.mark.parametrize("version", [114, 117, 118, 119, None])
    def test_read_empty_dta(self, version):
        empty_ds = DataFrame(columns=["unit"])
        # GH 7369, make sure can read a 0-obs dta file
        with tm.ensure_clean() as path:
            empty_ds.to_stata(path, write_index=False, version=version)
            empty_ds2 = read_stata(path)
            tm.assert_frame_equal(empty_ds, empty_ds2)

    @pytest.mark.parametrize("version", [114, 117, 118, 119, None])
    def test_read_empty_dta_with_dtypes(self, version):
        # GH 46240
        # Fixing above bug revealed that types are not correctly preserved when
        # writing empty DataFrames
        empty_df_typed = DataFrame(
            {
                "i8": np.array([0], dtype=np.int8),
                "i16": np.array([0], dtype=np.int16),
                "i32": np.array([0], dtype=np.int32),
                "i64": np.array([0], dtype=np.int64),
                "u8": np.array([0], dtype=np.uint8),
                "u16": np.array([0], dtype=np.uint16),
                "u32": np.array([0], dtype=np.uint32),
                "u64": np.array([0], dtype=np.uint64),
                "f32": np.array([0], dtype=np.float32),
                "f64": np.array([0], dtype=np.float64),
            }
        )
        expected = empty_df_typed.copy()
        # No uint# support. Downcast since values in range for int#
        expected["u8"] = expected["u8"].astype(np.int8)
        expected["u16"] = expected["u16"].astype(np.int16)
        expected["u32"] = expected["u32"].astype(np.int32)
        # No int64 supported at all. Downcast since values in range for int32
        expected["u64"] = expected["u64"].astype(np.int32)
        expected["i64"] = expected["i64"].astype(np.int32)

        # GH 7369, make sure can read a 0-obs dta file
        with tm.ensure_clean() as path:
            empty_df_typed.to_stata(path, write_index=False, version=version)
            empty_reread = read_stata(path)
            tm.assert_frame_equal(expected, empty_reread)
            tm.assert_series_equal(expected.dtypes, empty_reread.dtypes)

    @pytest.mark.parametrize("version", [114, 117, 118, 119, None])
    def test_read_index_col_none(self, version):
        df = DataFrame({"a": range(5), "b": ["b1", "b2", "b3", "b4", "b5"]})
        # GH 7369, make sure can read a 0-obs dta file
        with tm.ensure_clean() as path:
            df.to_stata(path, write_index=False, version=version)
            read_df = read_stata(path)

        assert isinstance(read_df.index, pd.RangeIndex)
        expected = df.copy()
        expected["a"] = expected["a"].astype(np.int32)
        tm.assert_frame_equal(read_df, expected, check_index_type=True)

    @pytest.mark.parametrize("file", ["stata1_114", "stata1_117"])
    def test_read_dta1(self, file, datapath):
        file = datapath("io", "data", "stata", f"{file}.dta")
        parsed = self.read_dta(file)

        # Pandas uses np.nan as missing value.
        # Thus, all columns will be of type float, regardless of their name.
        expected = DataFrame(
            [(np.nan, np.nan, np.nan, np.nan, np.nan)],
            columns=["float_miss", "double_miss", "byte_miss", "int_miss", "long_miss"],
        )

        # this is an oddity as really the nan should be float64, but
        # the casting doesn't fail so need to match stata here
        expected["float_miss"] = expected["float_miss"].astype(np.float32)

        tm.assert_frame_equal(parsed, expected)

    def test_read_dta2(self, datapath):
        expected = DataFrame.from_records(
            [
                (
                    datetime(2006, 11, 19, 23, 13, 20),
                    1479596223000,
                    datetime(2010, 1, 20),
                    datetime(2010, 1, 8),
                    datetime(2010, 1, 1),
                    datetime(1974, 7, 1),
                    datetime(2010, 1, 1),
                    datetime(2010, 1, 1),
                ),
                (
                    datetime(1959, 12, 31, 20, 3, 20),
                    -1479590,
                    datetime(1953, 10, 2),
                    datetime(1948, 6, 10),
                    datetime(1955, 1, 1),
                    datetime(1955, 7, 1),
                    datetime(1955, 1, 1),
                    datetime(2, 1, 1),
                ),
                (pd.NaT, pd.NaT, pd.NaT, pd.NaT, pd.NaT, pd.NaT, pd.NaT, pd.NaT),
            ],
            columns=[
                "datetime_c",
                "datetime_big_c",
                "date",
                "weekly_date",
                "monthly_date",
                "quarterly_date",
                "half_yearly_date",
                "yearly_date",
            ],
        )
        expected["yearly_date"] = expected["yearly_date"].astype("O")

        path1 = datapath("io", "data", "stata", "stata2_114.dta")
        path2 = datapath("io", "data", "stata", "stata2_115.dta")
        path3 = datapath("io", "data", "stata", "stata2_117.dta")

        with tm.assert_produces_warning(UserWarning):
            parsed_114 = self.read_dta(path1)
        with tm.assert_produces_warning(UserWarning):
            parsed_115 = self.read_dta(path2)
        with tm.assert_produces_warning(UserWarning):
            parsed_117 = self.read_dta(path3)
            # FIXME: don't leave commented-out
            # 113 is buggy due to limits of date format support in Stata
            # parsed_113 = self.read_dta(
            # datapath("io", "data", "stata", "stata2_113.dta")
            # )

        # FIXME: don't leave commented-out
        # buggy test because of the NaT comparison on certain platforms
        # Format 113 test fails since it does not support tc and tC formats
        # tm.assert_frame_equal(parsed_113, expected)
        tm.assert_frame_equal(parsed_114, expected, check_datetimelike_compat=True)
        tm.assert_frame_equal(parsed_115, expected, check_datetimelike_compat=True)
        tm.assert_frame_equal(parsed_117, expected, check_datetimelike_compat=True)

    @pytest.mark.parametrize(
        "file", ["stata3_113", "stata3_114", "stata3_115", "stata3_117"]
    )
    def test_read_dta3(self, file, datapath):
        file = datapath("io", "data", "stata", f"{file}.dta")
        parsed = self.read_dta(file)

        # match stata here
        expected = self.read_csv(datapath("io", "data", "stata", "stata3.csv"))
        expected = expected.astype(np.float32)
        expected["year"] = expected["year"].astype(np.int16)
        expected["quarter"] = expected["quarter"].astype(np.int8)

        tm.assert_frame_equal(parsed, expected)

    @pytest.mark.parametrize(
        "file", ["stata4_113", "stata4_114", "stata4_115", "stata4_117"]
    )
    def test_read_dta4(self, file, datapath):
        file = datapath("io", "data", "stata", f"{file}.dta")
        parsed = self.read_dta(file)

        expected = DataFrame.from_records(
            [
                ["one", "ten", "one", "one", "one"],
                ["two", "nine", "two", "two", "two"],
                ["three", "eight", "three", "three", "three"],
                ["four", "seven", 4, "four", "four"],
                ["five", "six", 5, np.nan, "five"],
                ["six", "five", 6, np.nan, "six"],
                ["seven", "four", 7, np.nan, "seven"],
                ["eight", "three", 8, np.nan, "eight"],
                ["nine", "two", 9, np.nan, "nine"],
                ["ten", "one", "ten", np.nan, "ten"],
            ],
            columns=[
                "fully_labeled",
                "fully_labeled2",
                "incompletely_labeled",
                "labeled_with_missings",
                "float_labelled",
            ],
        )

        # these are all categoricals
        for col in expected:
            orig = expected[col].copy()

            categories = np.asarray(expected["fully_labeled"][orig.notna()])
            if col == "incompletely_labeled":
                categories = orig

            cat = orig.astype("category")._values
            cat = cat.set_categories(categories, ordered=True)
            cat.categories.rename(None, inplace=True)

            expected[col] = cat

        # stata doesn't save .category metadata
        tm.assert_frame_equal(parsed, expected)

    # File containing strls
    def test_read_dta12(self, datapath):
        parsed_117 = self.read_dta(datapath("io", "data", "stata", "stata12_117.dta"))
        expected = DataFrame.from_records(
            [
                [1, "abc", "abcdefghi"],
                [3, "cba", "qwertywertyqwerty"],
                [93, "", "strl"],
            ],
            columns=["x", "y", "z"],
        )

        tm.assert_frame_equal(parsed_117, expected, check_dtype=False)

    def test_read_dta18(self, datapath):
        parsed_118 = self.read_dta(datapath("io", "data", "stata", "stata14_118.dta"))
        parsed_118["Bytes"] = parsed_118["Bytes"].astype("O")
        expected = DataFrame.from_records(
            [
                ["Cat", "Bogota", "Bogotá", 1, 1.0, "option b Ünicode", 1.0],
                ["Dog", "Boston", "Uzunköprü", np.nan, np.nan, np.nan, np.nan],
                ["Plane", "Rome", "Tromsø", 0, 0.0, "option a", 0.0],
                ["Potato", "Tokyo", "Elâzığ", -4, 4.0, 4, 4],  # noqa: RUF001
                ["", "", "", 0, 0.3332999, "option a", 1 / 3.0],
            ],
            columns=[
                "Things",
                "Cities",
                "Unicode_Cities_Strl",
                "Ints",
                "Floats",
                "Bytes",
                "Longs",
            ],
        )
        expected["Floats"] = expected["Floats"].astype(np.float32)
        for col in parsed_118.columns:
            tm.assert_almost_equal(parsed_118[col], expected[col])

        with StataReader(datapath("io", "data", "stata", "stata14_118.dta")) as rdr:
            vl = rdr.variable_labels()
            vl_expected = {
                "Unicode_Cities_Strl": "Here are some strls with Ünicode chars",
                "Longs": "long data",
                "Things": "Here are some things",
                "Bytes": "byte data",
                "Ints": "int data",
                "Cities": "Here are some cities",
                "Floats": "float data",
            }
            tm.assert_dict_equal(vl, vl_expected)

            assert rdr.data_label == "This is a  Ünicode data label"

    def test_read_write_dta5(self):
        original = DataFrame(
            [(np.nan, np.nan, np.nan, np.nan, np.nan)],
            columns=["float_miss", "double_miss", "byte_miss", "int_miss", "long_miss"],
        )
        original.index.name = "index"

        with tm.ensure_clean() as path:
            original.to_stata(path, convert_dates=None)
            written_and_read_again = self.read_dta(path)

        expected = original.copy()
        expected.index = expected.index.astype(np.int32)
        tm.assert_frame_equal(written_and_read_again.set_index("index"), expected)

    def test_write_dta6(self, datapath):
        original = self.read_csv(datapath("io", "data", "stata", "stata3.csv"))
        original.index.name = "index"
        original.index = original.index.astype(np.int32)
        original["year"] = original["year"].astype(np.int32)
        original["quarter"] = original["quarter"].astype(np.int32)

        with tm.ensure_clean() as path:
            original.to_stata(path, convert_dates=None)
            written_and_read_again = self.read_dta(path)
            tm.assert_frame_equal(
                written_and_read_again.set_index("index"),
                original,
                check_index_type=False,
            )

    @pytest.mark.parametrize("version", [114, 117, 118, 119, None])
    def test_read_write_dta10(self, version):
        original = DataFrame(
            data=[["string", "object", 1, 1.1, np.datetime64("2003-12-25")]],
            columns=["string", "object", "integer", "floating", "datetime"],
        )
        original["object"] = Series(original["object"], dtype=object)
        original.index.name = "index"
        original.index = original.index.astype(np.int32)
        original["integer"] = original["integer"].astype(np.int32)

        with tm.ensure_clean() as path:
            original.to_stata(path, convert_dates={"datetime": "tc"}, version=version)
            written_and_read_again = self.read_dta(path)
            # original.index is np.int32, read index is np.int64
            tm.assert_frame_equal(
                written_and_read_again.set_index("index"),
                original,
                check_index_type=False,
            )

    def test_stata_doc_examples(self):
        with tm.ensure_clean() as path:
            df = DataFrame(
                np.random.default_rng(2).standard_normal((10, 2)), columns=list("AB")
            )
            df.to_stata(path)

    def test_write_preserves_original(self):
        # 9795

        df = DataFrame(
            np.random.default_rng(2).standard_normal((5, 4)), columns=list("abcd")
        )
        df.loc[2, "a":"c"] = np.nan
        df_copy = df.copy()
        with tm.ensure_clean() as path:
            df.to_stata(path, write_index=False)
        tm.assert_frame_equal(df, df_copy)

    @pytest.mark.parametrize("version", [114, 117, 118, 119, None])
    def test_encoding(self, version, datapath):
        # GH 4626, proper encoding handling
        raw = read_stata(datapath("io", "data", "stata", "stata1_encoding.dta"))
        encoded = read_stata(datapath("io", "data", "stata", "stata1_encoding.dta"))
        result = encoded.kreis1849[0]

        expected = raw.kreis1849[0]
        assert result == expected
        assert isinstance(result, str)

        with tm.ensure_clean() as path:
            encoded.to_stata(path, write_index=False, version=version)
            reread_encoded = read_stata(path)
            tm.assert_frame_equal(encoded, reread_encoded)

    def test_read_write_dta11(self):
        original = DataFrame(
            [(1, 2, 3, 4)],
            columns=[
                "good",
                "b\u00E4d",
                "8number",
                "astringwithmorethan32characters______",
            ],
        )
        formatted = DataFrame(
            [(1, 2, 3, 4)],
            columns=["good", "b_d", "_8number", "astringwithmorethan32characters_"],
        )
        formatted.index.name = "index"
        formatted = formatted.astype(np.int32)

        with tm.ensure_clean() as path:
            with tm.assert_produces_warning(InvalidColumnName):
                original.to_stata(path, convert_dates=None)

            written_and_read_again = self.read_dta(path)

        expected = formatted.copy()
        expected.index = expected.index.astype(np.int32)
        tm.assert_frame_equal(written_and_read_again.set_index("index"), expected)

    @pytest.mark.parametrize("version", [114, 117, 118, 119, None])
    def test_read_write_dta12(self, version):
        original = DataFrame(
            [(1, 2, 3, 4, 5, 6)],
            columns=[
                "astringwithmorethan32characters_1",
                "astringwithmorethan32characters_2",
                "+",
                "-",
                "short",
                "delete",
            ],
        )
        formatted = DataFrame(
            [(1, 2, 3, 4, 5, 6)],
            columns=[
                "astringwithmorethan32characters_",
                "_0astringwithmorethan32character",
                "_",
                "_1_",
                "_short",
                "_delete",
            ],
        )
        formatted.index.name = "index"
        formatted = formatted.astype(np.int32)

        with tm.ensure_clean() as path:
            with tm.assert_produces_warning(InvalidColumnName):
                original.to_stata(path, convert_dates=None, version=version)
                # should get a warning for that format.

            written_and_read_again = self.read_dta(path)

        expected = formatted.copy()
        expected.index = expected.index.astype(np.int32)
        tm.assert_frame_equal(written_and_read_again.set_index("index"), expected)

    def test_read_write_dta13(self):
        s1 = Series(2**9, dtype=np.int16)
        s2 = Series(2**17, dtype=np.int32)
        s3 = Series(2**33, dtype=np.int64)
        original = DataFrame({"int16": s1, "int32": s2, "int64": s3})
        original.index.name = "index"

        formatted = original
        formatted["int64"] = formatted["int64"].astype(np.float64)

        with tm.ensure_clean() as path:
            original.to_stata(path)
            written_and_read_again = self.read_dta(path)

        expected = formatted.copy()
        expected.index = expected.index.astype(np.int32)
        tm.assert_frame_equal(written_and_read_again.set_index("index"), expected)

    @pytest.mark.parametrize("version", [114, 117, 118, 119, None])
    @pytest.mark.parametrize(
        "file", ["stata5_113", "stata5_114", "stata5_115", "stata5_117"]
    )
    def test_read_write_reread_dta14(self, file, parsed_114, version, datapath):
        file = datapath("io", "data", "stata", f"{file}.dta")
        parsed = self.read_dta(file)
        parsed.index.name = "index"

        tm.assert_frame_equal(parsed_114, parsed)

        with tm.ensure_clean() as path:
            parsed_114.to_stata(path, convert_dates={"date_td": "td"}, version=version)
            written_and_read_again = self.read_dta(path)

        expected = parsed_114.copy()
        expected.index = expected.index.astype(np.int32)
        tm.assert_frame_equal(written_and_read_again.set_index("index"), expected)

    @pytest.mark.parametrize(
        "file", ["stata6_113", "stata6_114", "stata6_115", "stata6_117"]
    )
    def test_read_write_reread_dta15(self, file, datapath):
        expected = self.read_csv(datapath("io", "data", "stata", "stata6.csv"))
        expected["byte_"] = expected["byte_"].astype(np.int8)
        expected["int_"] = expected["int_"].astype(np.int16)
        expected["long_"] = expected["long_"].astype(np.int32)
        expected["float_"] = expected["float_"].astype(np.float32)
        expected["double_"] = expected["double_"].astype(np.float64)
        expected["date_td"] = expected["date_td"].apply(
            datetime.strptime, args=("%Y-%m-%d",)
        )

        file = datapath("io", "data", "stata", f"{file}.dta")
        parsed = self.read_dta(file)

        tm.assert_frame_equal(expected, parsed)

    @pytest.mark.parametrize("version", [114, 117, 118, 119, None])
    def test_timestamp_and_label(self, version):
        original = DataFrame([(1,)], columns=["variable"])
        time_stamp = datetime(2000, 2, 29, 14, 21)
        data_label = "This is a data file."
        with tm.ensure_clean() as path:
            original.to_stata(
                path, time_stamp=time_stamp, data_label=data_label, version=version
            )

            with StataReader(path) as reader:
                assert reader.time_stamp == "29 Feb 2000 14:21"
                assert reader.data_label == data_label

    @pytest.mark.parametrize("version", [114, 117, 118, 119, None])
    def test_invalid_timestamp(self, version):
        original = DataFrame([(1,)], columns=["variable"])
        time_stamp = "01 Jan 2000, 00:00:00"
        with tm.ensure_clean() as path:
            msg = "time_stamp should be datetime type"
            with pytest.raises(ValueError, match=msg):
                original.to_stata(path, time_stamp=time_stamp, version=version)
            assert not os.path.isfile(path)

    def test_numeric_column_names(self):
        original = DataFrame(np.reshape(np.arange(25.0), (5, 5)))
        original.index.name = "index"
        with tm.ensure_clean() as path:
            # should get a warning for that format.
            with tm.assert_produces_warning(InvalidColumnName):
                original.to_stata(path)

            written_and_read_again = self.read_dta(path)

        written_and_read_again = written_and_read_again.set_index("index")
        columns = list(written_and_read_again.columns)
        convert_col_name = lambda x: int(x[1])
        written_and_read_again.columns = map(convert_col_name, columns)

        expected = original.copy()
        expected.index = expected.index.astype(np.int32)
        tm.assert_frame_equal(expected, written_and_read_again)

    @pytest.mark.parametrize("version", [114, 117, 118, 119, None])
    def test_nan_to_missing_value(self, version):
        s1 = Series(np.arange(4.0), dtype=np.float32)
        s2 = Series(np.arange(4.0), dtype=np.float64)
        s1[::2] = np.nan
        s2[1::2] = np.nan
        original = DataFrame({"s1": s1, "s2": s2})
        original.index.name = "index"

        with tm.ensure_clean() as path:
            original.to_stata(path, version=version)
            written_and_read_again = self.read_dta(path)

        written_and_read_again = written_and_read_again.set_index("index")
        expected = original.copy()
        expected.index = expected.index.astype(np.int32)
        tm.assert_frame_equal(written_and_read_again, expected)

    def test_no_index(self):
        columns = ["x", "y"]
        original = DataFrame(np.reshape(np.arange(10.0), (5, 2)), columns=columns)
        original.index.name = "index_not_written"
        with tm.ensure_clean() as path:
            original.to_stata(path, write_index=False)
            written_and_read_again = self.read_dta(path)
            with pytest.raises(KeyError, match=original.index.name):
                written_and_read_again["index_not_written"]

    def test_string_no_dates(self):
        s1 = Series(["a", "A longer string"])
        s2 = Series([1.0, 2.0], dtype=np.float64)
        original = DataFrame({"s1": s1, "s2": s2})
        original.index.name = "index"
        with tm.ensure_clean() as path:
            original.to_stata(path)
            written_and_read_again = self.read_dta(path)

        expected = original.copy()
        expected.index = expected.index.astype(np.int32)
        tm.assert_frame_equal(written_and_read_again.set_index("index"), expected)

    def test_large_value_conversion(self):
        s0 = Series([1, 99], dtype=np.int8)
        s1 = Series([1, 127], dtype=np.int8)
        s2 = Series([1, 2**15 - 1], dtype=np.int16)
        s3 = Series([1, 2**63 - 1], dtype=np.int64)
        original = DataFrame({"s0": s0, "s1": s1, "s2": s2, "s3": s3})
        original.index.name = "index"
        with tm.ensure_clean() as path:
            with tm.assert_produces_warning(PossiblePrecisionLoss):
                original.to_stata(path)

            written_and_read_again = self.read_dta(path)

        modified = original.copy()
        modified["s1"] = Series(modified["s1"], dtype=np.int16)
        modified["s2"] = Series(modified["s2"], dtype=np.int32)
        modified["s3"] = Series(modified["s3"], dtype=np.float64)
        modified.index = original.index.astype(np.int32)
        tm.assert_frame_equal(written_and_read_again.set_index("index"), modified)

    def test_dates_invalid_column(self):
        original = DataFrame([datetime(2006, 11, 19, 23, 13, 20)])
        original.index.name = "index"
        with tm.ensure_clean() as path:
            with tm.assert_produces_warning(InvalidColumnName):
                original.to_stata(path, convert_dates={0: "tc"})

            written_and_read_again = self.read_dta(path)

        modified = original.copy()
        modified.columns = ["_0"]
        modified.index = original.index.astype(np.int32)
        tm.assert_frame_equal(written_and_read_again.set_index("index"), modified)

    def test_105(self, datapath):
        # Data obtained from:
        # http://go.worldbank.org/ZXY29PVJ21
        dpath = datapath("io", "data", "stata", "S4_EDUC1.dta")
        df = read_stata(dpath)
        df0 = [[1, 1, 3, -2], [2, 1, 2, -2], [4, 1, 1, -2]]
        df0 = DataFrame(df0)
        df0.columns = ["clustnum", "pri_schl", "psch_num", "psch_dis"]
        df0["clustnum"] = df0["clustnum"].astype(np.int16)
        df0["pri_schl"] = df0["pri_schl"].astype(np.int8)
        df0["psch_num"] = df0["psch_num"].astype(np.int8)
        df0["psch_dis"] = df0["psch_dis"].astype(np.float32)
        tm.assert_frame_equal(df.head(3), df0)

    def test_value_labels_old_format(self, datapath):
        # GH 19417
        #
        # Test that value_labels() returns an empty dict if the file format
        # predates supporting value labels.
        dpath = datapath("io", "data", "stata", "S4_EDUC1.dta")
        with StataReader(dpath) as reader:
            assert reader.value_labels() == {}

    def test_date_export_formats(self):
        columns = ["tc", "td", "tw", "tm", "tq", "th", "ty"]
        conversions = {c: c for c in columns}
        data = [datetime(2006, 11, 20, 23, 13, 20)] * len(columns)
        original = DataFrame([data], columns=columns)
        original.index.name = "index"
        expected_values = [
            datetime(2006, 11, 20, 23, 13, 20),  # Time
            datetime(2006, 11, 20),  # Day
            datetime(2006, 11, 19),  # Week
            datetime(2006, 11, 1),  # Month
            datetime(2006, 10, 1),  # Quarter year
            datetime(2006, 7, 1),  # Half year
            datetime(2006, 1, 1),
        ]  # Year

        expected = DataFrame(
            [expected_values],
            index=pd.Index([0], dtype=np.int32, name="index"),
            columns=columns,
        )

        with tm.ensure_clean() as path:
            original.to_stata(path, convert_dates=conversions)
            written_and_read_again = self.read_dta(path)

        tm.assert_frame_equal(written_and_read_again.set_index("index"), expected)

    def test_write_missing_strings(self):
        original = DataFrame([["1"], [None]], columns=["foo"])

        expected = DataFrame(
            [["1"], [""]],
            index=pd.Index([0, 1], dtype=np.int32, name="index"),
            columns=["foo"],
        )

        with tm.ensure_clean() as path:
            original.to_stata(path)
            written_and_read_again = self.read_dta(path)

        tm.assert_frame_equal(written_and_read_again.set_index("index"), expected)

    @pytest.mark.parametrize("version", [114, 117, 118, 119, None])
    @pytest.mark.parametrize("byteorder", [">", "<"])
    def test_bool_uint(self, byteorder, version):
        s0 = Series([0, 1, True], dtype=np.bool_)
        s1 = Series([0, 1, 100], dtype=np.uint8)
        s2 = Series([0, 1, 255], dtype=np.uint8)
        s3 = Series([0, 1, 2**15 - 100], dtype=np.uint16)
        s4 = Series([0, 1, 2**16 - 1], dtype=np.uint16)
        s5 = Series([0, 1, 2**31 - 100], dtype=np.uint32)
        s6 = Series([0, 1, 2**32 - 1], dtype=np.uint32)

        original = DataFrame(
            {"s0": s0, "s1": s1, "s2": s2, "s3": s3, "s4": s4, "s5": s5, "s6": s6}
        )
        original.index.name = "index"
        expected = original.copy()
        expected.index = original.index.astype(np.int32)
        expected_types = (
            np.int8,
            np.int8,
            np.int16,
            np.int16,
            np.int32,
            np.int32,
            np.float64,
        )
        for c, t in zip(expected.columns, expected_types):
            expected[c] = expected[c].astype(t)

        with tm.ensure_clean() as path:
            original.to_stata(path, byteorder=byteorder, version=version)
            written_and_read_again = self.read_dta(path)

        written_and_read_again = written_and_read_again.set_index("index")
        tm.assert_frame_equal(written_and_read_again, expected)

    def test_variable_labels(self, datapath):
        with StataReader(datapath("io", "data", "stata", "stata7_115.dta")) as rdr:
            sr_115 = rdr.variable_labels()
        with StataReader(datapath("io", "data", "stata", "stata7_117.dta")) as rdr:
            sr_117 = rdr.variable_labels()
        keys = ("var1", "var2", "var3")
        labels = ("label1", "label2", "label3")
        for k, v in sr_115.items():
            assert k in sr_117
            assert v == sr_117[k]
            assert k in keys
            assert v in labels

    def test_minimal_size_col(self):
        str_lens = (1, 100, 244)
        s = {}
        for str_len in str_lens:
            s["s" + str(str_len)] = Series(
                ["a" * str_len, "b" * str_len, "c" * str_len]
            )
        original = DataFrame(s)
        with tm.ensure_clean() as path:
            original.to_stata(path, write_index=False)

            with StataReader(path) as sr:
                sr._ensure_open()  # The `_*list` variables are initialized here
                for variable, fmt, typ in zip(sr._varlist, sr._fmtlist, sr._typlist):
                    assert int(variable[1:]) == int(fmt[1:-1])
                    assert int(variable[1:]) == typ

    def test_excessively_long_string(self):
        str_lens = (1, 244, 500)
        s = {}
        for str_len in str_lens:
            s["s" + str(str_len)] = Series(
                ["a" * str_len, "b" * str_len, "c" * str_len]
            )
        original = DataFrame(s)
        msg = (
            r"Fixed width strings in Stata \.dta files are limited to 244 "
            r"\(or fewer\)\ncharacters\.  Column 's500' does not satisfy "
            r"this restriction\. Use the\n'version=117' parameter to write "
            r"the newer \(Stata 13 and later\) format\."
        )
        with pytest.raises(ValueError, match=msg):
            with tm.ensure_clean() as path:
                original.to_stata(path)

    def test_missing_value_generator(self):
        types = ("b", "h", "l")
        df = DataFrame([[0.0]], columns=["float_"])
        with tm.ensure_clean() as path:
            df.to_stata(path)
            with StataReader(path) as rdr:
                valid_range = rdr.VALID_RANGE
        expected_values = ["." + chr(97 + i) for i in range(26)]
        expected_values.insert(0, ".")
        for t in types:
            offset = valid_range[t][1]
            for i in range(27):
                val = StataMissingValue(offset + 1 + i)
                assert val.string == expected_values[i]

        # Test extremes for floats
        val = StataMissingValue(struct.unpack("<f", b"\x00\x00\x00\x7f")[0])
        assert val.string == "."
        val = StataMissingValue(struct.unpack("<f", b"\x00\xd0\x00\x7f")[0])
        assert val.string == ".z"

        # Test extremes for floats
        val = StataMissingValue(
            struct.unpack("<d", b"\x00\x00\x00\x00\x00\x00\xe0\x7f")[0]
        )
        assert val.string == "."
        val = StataMissingValue(
            struct.unpack("<d", b"\x00\x00\x00\x00\x00\x1a\xe0\x7f")[0]
        )
        assert val.string == ".z"

    @pytest.mark.parametrize("file", ["stata8_113", "stata8_115", "stata8_117"])
    def test_missing_value_conversion(self, file, datapath):
        columns = ["int8_", "int16_", "int32_", "float32_", "float64_"]
        smv = StataMissingValue(101)
        keys = sorted(smv.MISSING_VALUES.keys())
        data = []
        for i in range(27):
            row = [StataMissingValue(keys[i + (j * 27)]) for j in range(5)]
            data.append(row)
        expected = DataFrame(data, columns=columns)

        parsed = read_stata(
            datapath("io", "data", "stata", f"{file}.dta"), convert_missing=True
        )
        tm.assert_frame_equal(parsed, expected)

    def test_big_dates(self, datapath):
        yr = [1960, 2000, 9999, 100, 2262, 1677]
        mo = [1, 1, 12, 1, 4, 9]
        dd = [1, 1, 31, 1, 22, 23]
        hr = [0, 0, 23, 0, 0, 0]
        mm = [0, 0, 59, 0, 0, 0]
        ss = [0, 0, 59, 0, 0, 0]
        expected = []
        for year, month, day, hour, minute, second in zip(yr, mo, dd, hr, mm, ss):
            row = []
            for j in range(7):
                if j == 0:
                    row.append(datetime(year, month, day, hour, minute, second))
                elif j == 6:
                    row.append(datetime(year, 1, 1))
                else:
                    row.append(datetime(year, month, day))
            expected.append(row)
        expected.append([pd.NaT] * 7)
        columns = [
            "date_tc",
            "date_td",
            "date_tw",
            "date_tm",
            "date_tq",
            "date_th",
            "date_ty",
        ]

        # Fixes for weekly, quarterly,half,year
        expected[2][2] = datetime(9999, 12, 24)
        expected[2][3] = datetime(9999, 12, 1)
        expected[2][4] = datetime(9999, 10, 1)
        expected[2][5] = datetime(9999, 7, 1)
        expected[4][2] = datetime(2262, 4, 16)
        expected[4][3] = expected[4][4] = datetime(2262, 4, 1)
        expected[4][5] = expected[4][6] = datetime(2262, 1, 1)
        expected[5][2] = expected[5][3] = expected[5][4] = datetime(1677, 10, 1)
        expected[5][5] = expected[5][6] = datetime(1678, 1, 1)

        expected = DataFrame(expected, columns=columns, dtype=object)
        parsed_115 = read_stata(datapath("io", "data", "stata", "stata9_115.dta"))
        parsed_117 = read_stata(datapath("io", "data", "stata", "stata9_117.dta"))
        tm.assert_frame_equal(expected, parsed_115, check_datetimelike_compat=True)
        tm.assert_frame_equal(expected, parsed_117, check_datetimelike_compat=True)

        date_conversion = {c: c[-2:] for c in columns}
        # {c : c[-2:] for c in columns}
        with tm.ensure_clean() as path:
            expected.index.name = "index"
            expected.to_stata(path, convert_dates=date_conversion)
            written_and_read_again = self.read_dta(path)

        tm.assert_frame_equal(
            written_and_read_again.set_index("index"),
            expected.set_index(expected.index.astype(np.int32)),
            check_datetimelike_compat=True,
        )

    def test_dtype_conversion(self, datapath):
        expected = self.read_csv(datapath("io", "data", "stata", "stata6.csv"))
        expected["byte_"] = expected["byte_"].astype(np.int8)
        expected["int_"] = expected["int_"].astype(np.int16)
        expected["long_"] = expected["long_"].astype(np.int32)
        expected["float_"] = expected["float_"].astype(np.float32)
        expected["double_"] = expected["double_"].astype(np.float64)
        expected["date_td"] = expected["date_td"].apply(
            datetime.strptime, args=("%Y-%m-%d",)
        )

        no_conversion = read_stata(
            datapath("io", "data", "stata", "stata6_117.dta"), convert_dates=True
        )
        tm.assert_frame_equal(expected, no_conversion)

        conversion = read_stata(
            datapath("io", "data", "stata", "stata6_117.dta"),
            convert_dates=True,
            preserve_dtypes=False,
        )

        # read_csv types are the same
        expected = self.read_csv(datapath("io", "data", "stata", "stata6.csv"))
        expected["date_td"] = expected["date_td"].apply(
            datetime.strptime, args=("%Y-%m-%d",)
        )

        tm.assert_frame_equal(expected, conversion)

    def test_drop_column(self, datapath):
        expected = self.read_csv(datapath("io", "data", "stata", "stata6.csv"))
        expected["byte_"] = expected["byte_"].astype(np.int8)
        expected["int_"] = expected["int_"].astype(np.int16)
        expected["long_"] = expected["long_"].astype(np.int32)
        expected["float_"] = expected["float_"].astype(np.float32)
        expected["double_"] = expected["double_"].astype(np.float64)
        expected["date_td"] = expected["date_td"].apply(
            datetime.strptime, args=("%Y-%m-%d",)
        )

        columns = ["byte_", "int_", "long_"]
        expected = expected[columns]
        dropped = read_stata(
            datapath("io", "data", "stata", "stata6_117.dta"),
            convert_dates=True,
            columns=columns,
        )

        tm.assert_frame_equal(expected, dropped)

        # See PR 10757
        columns = ["int_", "long_", "byte_"]
        expected = expected[columns]
        reordered = read_stata(
            datapath("io", "data", "stata", "stata6_117.dta"),
            convert_dates=True,
            columns=columns,
        )
        tm.assert_frame_equal(expected, reordered)

        msg = "columns contains duplicate entries"
        with pytest.raises(ValueError, match=msg):
            columns = ["byte_", "byte_"]
            read_stata(
                datapath("io", "data", "stata", "stata6_117.dta"),
                convert_dates=True,
                columns=columns,
            )

        msg = "The following columns were not found in the Stata data set: not_found"
        with pytest.raises(ValueError, match=msg):
            columns = ["byte_", "int_", "long_", "not_found"]
            read_stata(
                datapath("io", "data", "stata", "stata6_117.dta"),
                convert_dates=True,
                columns=columns,
            )

    @pytest.mark.parametrize("version", [114, 117, 118, 119, None])
    @pytest.mark.filterwarnings(
        "ignore:\\nStata value:pandas.io.stata.ValueLabelTypeMismatch"
    )
    def test_categorical_writing(self, version):
        original = DataFrame.from_records(
            [
                ["one", "ten", "one", "one", "one", 1],
                ["two", "nine", "two", "two", "two", 2],
                ["three", "eight", "three", "three", "three", 3],
                ["four", "seven", 4, "four", "four", 4],
                ["five", "six", 5, np.nan, "five", 5],
                ["six", "five", 6, np.nan, "six", 6],
                ["seven", "four", 7, np.nan, "seven", 7],
                ["eight", "three", 8, np.nan, "eight", 8],
                ["nine", "two", 9, np.nan, "nine", 9],
                ["ten", "one", "ten", np.nan, "ten", 10],
            ],
            columns=[
                "fully_labeled",
                "fully_labeled2",
                "incompletely_labeled",
                "labeled_with_missings",
                "float_labelled",
                "unlabeled",
            ],
        )
        expected = original.copy()

        # these are all categoricals
        original = pd.concat(
            [original[col].astype("category") for col in original], axis=1
        )
        expected.index = expected.index.set_names("index").astype(np.int32)

        expected["incompletely_labeled"] = expected["incompletely_labeled"].apply(str)
        expected["unlabeled"] = expected["unlabeled"].apply(str)
        for col in expected:
            orig = expected[col].copy()

            cat = orig.astype("category")._values
            cat = cat.as_ordered()
            if col == "unlabeled":
                cat = cat.set_categories(orig, ordered=True)

            cat.categories.rename(None, inplace=True)

            expected[col] = cat

        with tm.ensure_clean() as path:
            original.to_stata(path, version=version)
            written_and_read_again = self.read_dta(path)

        res = written_and_read_again.set_index("index")
        tm.assert_frame_equal(res, expected)

    def test_categorical_warnings_and_errors(self):
        # Warning for non-string labels
        # Error for labels too long
        original = DataFrame.from_records(
            [["a" * 10000], ["b" * 10000], ["c" * 10000], ["d" * 10000]],
            columns=["Too_long"],
        )

        original = pd.concat(
            [original[col].astype("category") for col in original], axis=1
        )
        with tm.ensure_clean() as path:
            msg = (
                "Stata value labels for a single variable must have "
                r"a combined length less than 32,000 characters\."
            )
            with pytest.raises(ValueError, match=msg):
                original.to_stata(path)

        original = DataFrame.from_records(
            [["a"], ["b"], ["c"], ["d"], [1]], columns=["Too_long"]
        )
        original = pd.concat(
            [original[col].astype("category") for col in original], axis=1
        )

        with tm.assert_produces_warning(ValueLabelTypeMismatch):
            original.to_stata(path)
            # should get a warning for mixed content

    @pytest.mark.parametrize("version", [114, 117, 118, 119, None])
    def test_categorical_with_stata_missing_values(self, version):
        values = [["a" + str(i)] for i in range(120)]
        values.append([np.nan])
        original = DataFrame.from_records(values, columns=["many_labels"])
        original = pd.concat(
            [original[col].astype("category") for col in original], axis=1
        )
        original.index.name = "index"
        with tm.ensure_clean() as path:
            original.to_stata(path, version=version)
            written_and_read_again = self.read_dta(path)

        res = written_and_read_again.set_index("index")

        expected = original.copy()
        for col in expected:
            cat = expected[col]._values
            new_cats = cat.remove_unused_categories().categories
            cat = cat.set_categories(new_cats, ordered=True)
            expected[col] = cat
        expected.index = expected.index.astype(np.int32)
        tm.assert_frame_equal(res, expected)

    @pytest.mark.parametrize("file", ["stata10_115", "stata10_117"])
    def test_categorical_order(self, file, datapath):
        # Directly construct using expected codes
        # Format is is_cat, col_name, labels (in order), underlying data
        expected = [
            (True, "ordered", ["a", "b", "c", "d", "e"], np.arange(5)),
            (True, "reverse", ["a", "b", "c", "d", "e"], np.arange(5)[::-1]),
            (True, "noorder", ["a", "b", "c", "d", "e"], np.array([2, 1, 4, 0, 3])),
            (True, "floating", ["a", "b", "c", "d", "e"], np.arange(0, 5)),
            (True, "float_missing", ["a", "d", "e"], np.array([0, 1, 2, -1, -1])),
            (False, "nolabel", [1.0, 2.0, 3.0, 4.0, 5.0], np.arange(5)),
            (True, "int32_mixed", ["d", 2, "e", "b", "a"], np.arange(5)),
        ]
        cols = []
        for is_cat, col, labels, codes in expected:
            if is_cat:
                cols.append(
                    (col, pd.Categorical.from_codes(codes, labels, ordered=True))
                )
            else:
                cols.append((col, Series(labels, dtype=np.float32)))
        expected = DataFrame.from_dict(dict(cols))

        # Read with and with out categoricals, ensure order is identical
        file = datapath("io", "data", "stata", f"{file}.dta")
        parsed = read_stata(file)
        tm.assert_frame_equal(expected, parsed)

        # Check identity of codes
        for col in expected:
            if isinstance(expected[col].dtype, CategoricalDtype):
                tm.assert_series_equal(expected[col].cat.codes, parsed[col].cat.codes)
                tm.assert_index_equal(
                    expected[col].cat.categories, parsed[col].cat.categories
                )

    @pytest.mark.parametrize("file", ["stata11_115", "stata11_117"])
    def test_categorical_sorting(self, file, datapath):
        parsed = read_stata(datapath("io", "data", "stata", f"{file}.dta"))

        # Sort based on codes, not strings
        parsed = parsed.sort_values("srh", na_position="first")

        # Don't sort index
        parsed.index = pd.RangeIndex(len(parsed))
        codes = [-1, -1, 0, 1, 1, 1, 2, 2, 3, 4]
        categories = ["Poor", "Fair", "Good", "Very good", "Excellent"]
        cat = pd.Categorical.from_codes(
            codes=codes, categories=categories, ordered=True
        )
        expected = Series(cat, name="srh")
        tm.assert_series_equal(expected, parsed["srh"])

    @pytest.mark.parametrize("file", ["stata10_115", "stata10_117"])
    def test_categorical_ordering(self, file, datapath):
        file = datapath("io", "data", "stata", f"{file}.dta")
        parsed = read_stata(file)

        parsed_unordered = read_stata(file, order_categoricals=False)
        for col in parsed:
            if not isinstance(parsed[col].dtype, CategoricalDtype):
                continue
            assert parsed[col].cat.ordered
            assert not parsed_unordered[col].cat.ordered

    @pytest.mark.filterwarnings("ignore::UserWarning")
    @pytest.mark.parametrize(
        "file",
        [
            "stata1_117",
            "stata2_117",
            "stata3_117",
            "stata4_117",
            "stata5_117",
            "stata6_117",
            "stata7_117",
            "stata8_117",
            "stata9_117",
            "stata10_117",
            "stata11_117",
        ],
    )
    @pytest.mark.parametrize("chunksize", [1, 2])
    @pytest.mark.parametrize("convert_categoricals", [False, True])
    @pytest.mark.parametrize("convert_dates", [False, True])
    def test_read_chunks_117(
        self, file, chunksize, convert_categoricals, convert_dates, datapath
    ):
        fname = datapath("io", "data", "stata", f"{file}.dta")

        parsed = read_stata(
            fname,
            convert_categoricals=convert_categoricals,
            convert_dates=convert_dates,
        )
        with read_stata(
            fname,
            iterator=True,
            convert_categoricals=convert_categoricals,
            convert_dates=convert_dates,
        ) as itr:
            pos = 0
            for j in range(5):
                try:
                    chunk = itr.read(chunksize)
                except StopIteration:
                    break
                from_frame = parsed.iloc[pos : pos + chunksize, :].copy()
                from_frame = self._convert_categorical(from_frame)
                tm.assert_frame_equal(
                    from_frame, chunk, check_dtype=False, check_datetimelike_compat=True
                )
                pos += chunksize

    @staticmethod
    def _convert_categorical(from_frame: DataFrame) -> DataFrame:
        """
        Emulate the categorical casting behavior we expect from roundtripping.
        """
        for col in from_frame:
            ser = from_frame[col]
            if isinstance(ser.dtype, CategoricalDtype):
                cat = ser._values.remove_unused_categories()
                if cat.categories.dtype == object:
                    categories = pd.Index._with_infer(cat.categories._values)
                    cat = cat.set_categories(categories)
                from_frame[col] = cat
        return from_frame

    def test_iterator(self, datapath):
        fname = datapath("io", "data", "stata", "stata3_117.dta")

        parsed = read_stata(fname)

        with read_stata(fname, iterator=True) as itr:
            chunk = itr.read(5)
            tm.assert_frame_equal(parsed.iloc[0:5, :], chunk)

        with read_stata(fname, chunksize=5) as itr:
            chunk = list(itr)
            tm.assert_frame_equal(parsed.iloc[0:5, :], chunk[0])

        with read_stata(fname, iterator=True) as itr:
            chunk = itr.get_chunk(5)
            tm.assert_frame_equal(parsed.iloc[0:5, :], chunk)

        with read_stata(fname, chunksize=5) as itr:
            chunk = itr.get_chunk()
            tm.assert_frame_equal(parsed.iloc[0:5, :], chunk)

        # GH12153
        with read_stata(fname, chunksize=4) as itr:
            from_chunks = pd.concat(itr)
        tm.assert_frame_equal(parsed, from_chunks)

    @pytest.mark.filterwarnings("ignore::UserWarning")
    @pytest.mark.parametrize(
        "file",
        [
            "stata2_115",
            "stata3_115",
            "stata4_115",
            "stata5_115",
            "stata6_115",
            "stata7_115",
            "stata8_115",
            "stata9_115",
            "stata10_115",
            "stata11_115",
        ],
    )
    @pytest.mark.parametrize("chunksize", [1, 2])
    @pytest.mark.parametrize("convert_categoricals", [False, True])
    @pytest.mark.parametrize("convert_dates", [False, True])
    def test_read_chunks_115(
        self, file, chunksize, convert_categoricals, convert_dates, datapath
    ):
        fname = datapath("io", "data", "stata", f"{file}.dta")

        # Read the whole file
        parsed = read_stata(
            fname,
            convert_categoricals=convert_categoricals,
            convert_dates=convert_dates,
        )

        # Compare to what we get when reading by chunk
        with read_stata(
            fname,
            iterator=True,
            convert_dates=convert_dates,
            convert_categoricals=convert_categoricals,
        ) as itr:
            pos = 0
            for j in range(5):
                try:
                    chunk = itr.read(chunksize)
                except StopIteration:
                    break
                from_frame = parsed.iloc[pos : pos + chunksize, :].copy()
                from_frame = self._convert_categorical(from_frame)
                tm.assert_frame_equal(
                    from_frame, chunk, check_dtype=False, check_datetimelike_compat=True
                )
                pos += chunksize

    def test_read_chunks_columns(self, datapath):
        fname = datapath("io", "data", "stata", "stata3_117.dta")
        columns = ["quarter", "cpi", "m1"]
        chunksize = 2

        parsed = read_stata(fname, columns=columns)
        with read_stata(fname, iterator=True) as itr:
            pos = 0
            for j in range(5):
                chunk = itr.read(chunksize, columns=columns)
                if chunk is None:
                    break
                from_frame = parsed.iloc[pos : pos + chunksize, :]
                tm.assert_frame_equal(from_frame, chunk, check_dtype=False)
                pos += chunksize

    @pytest.mark.parametrize("version", [114, 117, 118, 119, None])
    def test_write_variable_labels(self, version, mixed_frame):
        # GH 13631, add support for writing variable labels
        mixed_frame.index.name = "index"
        variable_labels = {"a": "City Rank", "b": "City Exponent", "c": "City"}
        with tm.ensure_clean() as path:
            mixed_frame.to_stata(path, variable_labels=variable_labels, version=version)
            with StataReader(path) as sr:
                read_labels = sr.variable_labels()
            expected_labels = {
                "index": "",
                "a": "City Rank",
                "b": "City Exponent",
                "c": "City",
            }
            assert read_labels == expected_labels

        variable_labels["index"] = "The Index"
        with tm.ensure_clean() as path:
            mixed_frame.to_stata(path, variable_labels=variable_labels, version=version)
            with StataReader(path) as sr:
                read_labels = sr.variable_labels()
            assert read_labels == variable_labels

    @pytest.mark.parametrize("version", [114, 117, 118, 119, None])
    def test_invalid_variable_labels(self, version, mixed_frame):
        mixed_frame.index.name = "index"
        variable_labels = {"a": "very long" * 10, "b": "City Exponent", "c": "City"}
        with tm.ensure_clean() as path:
            msg = "Variable labels must be 80 characters or fewer"
            with pytest.raises(ValueError, match=msg):
                mixed_frame.to_stata(
                    path, variable_labels=variable_labels, version=version
                )

    @pytest.mark.parametrize("version", [114, 117])
    def test_invalid_variable_label_encoding(self, version, mixed_frame):
        mixed_frame.index.name = "index"
        variable_labels = {"a": "very long" * 10, "b": "City Exponent", "c": "City"}
        variable_labels["a"] = "invalid character Œ"
        with tm.ensure_clean() as path:
            with pytest.raises(
                ValueError, match="Variable labels must contain only characters"
            ):
                mixed_frame.to_stata(
                    path, variable_labels=variable_labels, version=version
                )

    def test_write_variable_label_errors(self, mixed_frame):
        values = ["\u03A1", "\u0391", "\u039D", "\u0394", "\u0391", "\u03A3"]

        variable_labels_utf8 = {
            "a": "City Rank",
            "b": "City Exponent",
            "c": "".join(values),
        }

        msg = (
            "Variable labels must contain only characters that can be "
            "encoded in Latin-1"
        )
        with pytest.raises(ValueError, match=msg):
            with tm.ensure_clean() as path:
                mixed_frame.to_stata(path, variable_labels=variable_labels_utf8)

        variable_labels_long = {
            "a": "City Rank",
            "b": "City Exponent",
            "c": "A very, very, very long variable label "
            "that is too long for Stata which means "
            "that it has more than 80 characters",
        }

        msg = "Variable labels must be 80 characters or fewer"
        with pytest.raises(ValueError, match=msg):
            with tm.ensure_clean() as path:
                mixed_frame.to_stata(path, variable_labels=variable_labels_long)

    def test_default_date_conversion(self):
        # GH 12259
        dates = [
            dt.datetime(1999, 12, 31, 12, 12, 12, 12000),
            dt.datetime(2012, 12, 21, 12, 21, 12, 21000),
            dt.datetime(1776, 7, 4, 7, 4, 7, 4000),
        ]
        original = DataFrame(
            {
                "nums": [1.0, 2.0, 3.0],
                "strs": ["apple", "banana", "cherry"],
                "dates": dates,
            }
        )

        with tm.ensure_clean() as path:
            original.to_stata(path, write_index=False)
            reread = read_stata(path, convert_dates=True)
            tm.assert_frame_equal(original, reread)

            original.to_stata(path, write_index=False, convert_dates={"dates": "tc"})
            direct = read_stata(path, convert_dates=True)
            tm.assert_frame_equal(reread, direct)

            dates_idx = original.columns.tolist().index("dates")
            original.to_stata(path, write_index=False, convert_dates={dates_idx: "tc"})
            direct = read_stata(path, convert_dates=True)
            tm.assert_frame_equal(reread, direct)

    def test_unsupported_type(self):
        original = DataFrame({"a": [1 + 2j, 2 + 4j]})

        msg = "Data type complex128 not supported"
        with pytest.raises(NotImplementedError, match=msg):
            with tm.ensure_clean() as path:
                original.to_stata(path)

    def test_unsupported_datetype(self):
        dates = [
            dt.datetime(1999, 12, 31, 12, 12, 12, 12000),
            dt.datetime(2012, 12, 21, 12, 21, 12, 21000),
            dt.datetime(1776, 7, 4, 7, 4, 7, 4000),
        ]
        original = DataFrame(
            {
                "nums": [1.0, 2.0, 3.0],
                "strs": ["apple", "banana", "cherry"],
                "dates": dates,
            }
        )

        msg = "Format %tC not implemented"
        with pytest.raises(NotImplementedError, match=msg):
            with tm.ensure_clean() as path:
                original.to_stata(path, convert_dates={"dates": "tC"})

        dates = pd.date_range("1-1-1990", periods=3, tz="Asia/Hong_Kong")
        original = DataFrame(
            {
                "nums": [1.0, 2.0, 3.0],
                "strs": ["apple", "banana", "cherry"],
                "dates": dates,
            }
        )
        with pytest.raises(NotImplementedError, match="Data type datetime64"):
            with tm.ensure_clean() as path:
                original.to_stata(path)

    def test_repeated_column_labels(self, datapath):
        # GH 13923, 25772
        msg = """
Value labels for column ethnicsn are not unique. These cannot be converted to
pandas categoricals.

Either read the file with `convert_categoricals` set to False or use the
low level interface in `StataReader` to separately read the values and the
value_labels.

The repeated labels are:\n-+\nwolof
"""
        with pytest.raises(ValueError, match=msg):
            read_stata(
                datapath("io", "data", "stata", "stata15.dta"),
                convert_categoricals=True,
            )

    def test_stata_111(self, datapath):
        # 111 is an old version but still used by current versions of
        # SAS when exporting to Stata format. We do not know of any
        # on-line documentation for this version.
        df = read_stata(datapath("io", "data", "stata", "stata7_111.dta"))
        original = DataFrame(
            {
                "y": [1, 1, 1, 1, 1, 0, 0, np.nan, 0, 0],
                "x": [1, 2, 1, 3, np.nan, 4, 3, 5, 1, 6],
                "w": [2, np.nan, 5, 2, 4, 4, 3, 1, 2, 3],
                "z": ["a", "b", "c", "d", "e", "", "g", "h", "i", "j"],
            }
        )
        original = original[["y", "x", "w", "z"]]
        tm.assert_frame_equal(original, df)

    def test_out_of_range_double(self):
        # GH 14618
        df = DataFrame(
            {
                "ColumnOk": [0.0, np.finfo(np.double).eps, 4.49423283715579e307],
                "ColumnTooBig": [0.0, np.finfo(np.double).eps, np.finfo(np.double).max],
            }
        )
        msg = (
            r"Column ColumnTooBig has a maximum value \(.+\) outside the range "
            r"supported by Stata \(.+\)"
        )
        with pytest.raises(ValueError, match=msg):
            with tm.ensure_clean() as path:
                df.to_stata(path)

    def test_out_of_range_float(self):
        original = DataFrame(
            {
                "ColumnOk": [
                    0.0,
                    np.finfo(np.float32).eps,
                    np.finfo(np.float32).max / 10.0,
                ],
                "ColumnTooBig": [
                    0.0,
                    np.finfo(np.float32).eps,
                    np.finfo(np.float32).max,
                ],
            }
        )
        original.index.name = "index"
        for col in original:
            original[col] = original[col].astype(np.float32)

        with tm.ensure_clean() as path:
            original.to_stata(path)
            reread = read_stata(path)

        original["ColumnTooBig"] = original["ColumnTooBig"].astype(np.float64)
        expected = original.copy()
        expected.index = expected.index.astype(np.int32)
        tm.assert_frame_equal(reread.set_index("index"), expected)

    @pytest.mark.parametrize("infval", [np.inf, -np.inf])
    def test_inf(self, infval):
        # GH 45350
        df = DataFrame({"WithoutInf": [0.0, 1.0], "WithInf": [2.0, infval]})
        msg = (
            "Column WithInf contains infinity or -infinity"
            "which is outside the range supported by Stata."
        )
        with pytest.raises(ValueError, match=msg):
            with tm.ensure_clean() as path:
                df.to_stata(path)

    def test_path_pathlib(self):
        df = DataFrame(
            1.1 * np.arange(120).reshape((30, 4)),
            columns=pd.Index(list("ABCD"), dtype=object),
            index=pd.Index([f"i-{i}" for i in range(30)], dtype=object),
        )
        df.index.name = "index"
        reader = lambda x: read_stata(x).set_index("index")
        result = tm.round_trip_pathlib(df.to_stata, reader)
        tm.assert_frame_equal(df, result)

    def test_pickle_path_localpath(self):
        df = DataFrame(
            1.1 * np.arange(120).reshape((30, 4)),
            columns=pd.Index(list("ABCD"), dtype=object),
            index=pd.Index([f"i-{i}" for i in range(30)], dtype=object),
        )
        df.index.name = "index"
        reader = lambda x: read_stata(x).set_index("index")
        result = tm.round_trip_localpath(df.to_stata, reader)
        tm.assert_frame_equal(df, result)

    @pytest.mark.parametrize("write_index", [True, False])
    def test_value_labels_iterator(self, write_index):
        # GH 16923
        d = {"A": ["B", "E", "C", "A", "E"]}
        df = DataFrame(data=d)
        df["A"] = df["A"].astype("category")
        with tm.ensure_clean() as path:
            df.to_stata(path, write_index=write_index)

            with read_stata(path, iterator=True) as dta_iter:
                value_labels = dta_iter.value_labels()
        assert value_labels == {"A": {0: "A", 1: "B", 2: "C", 3: "E"}}

    def test_set_index(self):
        # GH 17328
        df = DataFrame(
            1.1 * np.arange(120).reshape((30, 4)),
            columns=pd.Index(list("ABCD"), dtype=object),
            index=pd.Index([f"i-{i}" for i in range(30)], dtype=object),
        )
        df.index.name = "index"
        with tm.ensure_clean() as path:
            df.to_stata(path)
            reread = read_stata(path, index_col="index")
        tm.assert_frame_equal(df, reread)

    @pytest.mark.parametrize(
        "column", ["ms", "day", "week", "month", "qtr", "half", "yr"]
    )
    def test_date_parsing_ignores_format_details(self, column, datapath):
        # GH 17797
        #
        # Test that display formats are ignored when determining if a numeric
        # column is a date value.
        #
        # All date types are stored as numbers and format associated with the
        # column denotes both the type of the date and the display format.
        #
        # STATA supports 9 date types which each have distinct units. We test 7
        # of the 9 types, ignoring %tC and %tb. %tC is a variant of %tc that
        # accounts for leap seconds and %tb relies on STATAs business calendar.
        df = read_stata(datapath("io", "data", "stata", "stata13_dates.dta"))
        unformatted = df.loc[0, column]
        formatted = df.loc[0, column + "_fmt"]
        assert unformatted == formatted

    def test_writer_117(self):
        original = DataFrame(
            data=[
                [
                    "string",
                    "object",
                    1,
                    1,
                    1,
                    1.1,
                    1.1,
                    np.datetime64("2003-12-25"),
                    "a",
                    "a" * 2045,
                    "a" * 5000,
                    "a",
                ],
                [
                    "string-1",
                    "object-1",
                    1,
                    1,
                    1,
                    1.1,
                    1.1,
                    np.datetime64("2003-12-26"),
                    "b",
                    "b" * 2045,
                    "",
                    "",
                ],
            ],
            columns=[
                "string",
                "object",
                "int8",
                "int16",
                "int32",
                "float32",
                "float64",
                "datetime",
                "s1",
                "s2045",
                "srtl",
                "forced_strl",
            ],
        )
        original["object"] = Series(original["object"], dtype=object)
        original["int8"] = Series(original["int8"], dtype=np.int8)
        original["int16"] = Series(original["int16"], dtype=np.int16)
        original["int32"] = original["int32"].astype(np.int32)
        original["float32"] = Series(original["float32"], dtype=np.float32)
        original.index.name = "index"
        original.index = original.index.astype(np.int32)
        copy = original.copy()
        with tm.ensure_clean() as path:
            original.to_stata(
                path,
                convert_dates={"datetime": "tc"},
                convert_strl=["forced_strl"],
                version=117,
            )
            written_and_read_again = self.read_dta(path)
            # original.index is np.int32, read index is np.int64
            tm.assert_frame_equal(
                written_and_read_again.set_index("index"),
                original,
                check_index_type=False,
            )
            tm.assert_frame_equal(original, copy)

    def test_convert_strl_name_swap(self):
        original = DataFrame(
            [["a" * 3000, "A", "apple"], ["b" * 1000, "B", "banana"]],
            columns=["long1" * 10, "long", 1],
        )
        original.index.name = "index"

        with tm.assert_produces_warning(InvalidColumnName):
            with tm.ensure_clean() as path:
                original.to_stata(path, convert_strl=["long", 1], version=117)
                reread = self.read_dta(path)
                reread = reread.set_index("index")
                reread.columns = original.columns
                tm.assert_frame_equal(reread, original, check_index_type=False)

    def test_invalid_date_conversion(self):
        # GH 12259
        dates = [
            dt.datetime(1999, 12, 31, 12, 12, 12, 12000),
            dt.datetime(2012, 12, 21, 12, 21, 12, 21000),
            dt.datetime(1776, 7, 4, 7, 4, 7, 4000),
        ]
        original = DataFrame(
            {
                "nums": [1.0, 2.0, 3.0],
                "strs": ["apple", "banana", "cherry"],
                "dates": dates,
            }
        )

        with tm.ensure_clean() as path:
            msg = "convert_dates key must be a column or an integer"
            with pytest.raises(ValueError, match=msg):
                original.to_stata(path, convert_dates={"wrong_name": "tc"})

    @pytest.mark.parametrize("version", [114, 117, 118, 119, None])
    def test_nonfile_writing(self, version):
        # GH 21041
        bio = io.BytesIO()
        df = DataFrame(
            1.1 * np.arange(120).reshape((30, 4)),
            columns=pd.Index(list("ABCD"), dtype=object),
            index=pd.Index([f"i-{i}" for i in range(30)], dtype=object),
        )
        df.index.name = "index"
        with tm.ensure_clean() as path:
            df.to_stata(bio, version=version)
            bio.seek(0)
            with open(path, "wb") as dta:
                dta.write(bio.read())
            reread = read_stata(path, index_col="index")
        tm.assert_frame_equal(df, reread)

    def test_gzip_writing(self):
        # writing version 117 requires seek and cannot be used with gzip
        df = DataFrame(
            1.1 * np.arange(120).reshape((30, 4)),
            columns=pd.Index(list("ABCD"), dtype=object),
            index=pd.Index([f"i-{i}" for i in range(30)], dtype=object),
        )
        df.index.name = "index"
        with tm.ensure_clean() as path:
            with gzip.GzipFile(path, "wb") as gz:
                df.to_stata(gz, version=114)
            with gzip.GzipFile(path, "rb") as gz:
                reread = read_stata(gz, index_col="index")
        tm.assert_frame_equal(df, reread)

    def test_unicode_dta_118(self, datapath):
        unicode_df = self.read_dta(datapath("io", "data", "stata", "stata16_118.dta"))

        columns = ["utf8", "latin1", "ascii", "utf8_strl", "ascii_strl"]
        values = [
            ["ραηδας", "PÄNDÄS", "p", "ραηδας", "p"],
            ["ƤĀńĐąŜ", "Ö", "a", "ƤĀńĐąŜ", "a"],
            ["ᴘᴀᴎᴅᴀS", "Ü", "n", "ᴘᴀᴎᴅᴀS", "n"],
            ["      ", "      ", "d", "      ", "d"],
            [" ", "", "a", " ", "a"],
            ["", "", "s", "", "s"],
            ["", "", " ", "", " "],
        ]
        expected = DataFrame(values, columns=columns)

        tm.assert_frame_equal(unicode_df, expected)

    def test_mixed_string_strl(self):
        # GH 23633
        output = [{"mixed": "string" * 500, "number": 0}, {"mixed": None, "number": 1}]
        output = DataFrame(output)
        output.number = output.number.astype("int32")

        with tm.ensure_clean() as path:
            output.to_stata(path, write_index=False, version=117)
            reread = read_stata(path)
            expected = output.fillna("")
            tm.assert_frame_equal(reread, expected)

            # Check strl supports all None (null)
            output["mixed"] = None
            output.to_stata(
                path, write_index=False, convert_strl=["mixed"], version=117
            )
            reread = read_stata(path)
            expected = output.fillna("")
            tm.assert_frame_equal(reread, expected)

    @pytest.mark.parametrize("version", [114, 117, 118, 119, None])
    def test_all_none_exception(self, version):
        output = [{"none": "none", "number": 0}, {"none": None, "number": 1}]
        output = DataFrame(output)
        output["none"] = None
        with tm.ensure_clean() as path:
            with pytest.raises(ValueError, match="Column `none` cannot be exported"):
                output.to_stata(path, version=version)

    @pytest.mark.parametrize("version", [114, 117, 118, 119, None])
    def test_invalid_file_not_written(self, version):
        content = "Here is one __�__ Another one __·__ Another one __½__"
        df = DataFrame([content], columns=["invalid"])
        with tm.ensure_clean() as path:
            msg1 = (
                r"'latin-1' codec can't encode character '\\ufffd' "
                r"in position 14: ordinal not in range\(256\)"
            )
            msg2 = (
                "'ascii' codec can't decode byte 0xef in position 14: "
                r"ordinal not in range\(128\)"
            )
            with pytest.raises(UnicodeEncodeError, match=f"{msg1}|{msg2}"):
                df.to_stata(path)

    def test_strl_latin1(self):
        # GH 23573, correct GSO data to reflect correct size
        output = DataFrame(
            [["pandas"] * 2, ["þâÑÐŧ"] * 2], columns=["var_str", "var_strl"]
        )

        with tm.ensure_clean() as path:
            output.to_stata(path, version=117, convert_strl=["var_strl"])
            with open(path, "rb") as reread:
                content = reread.read()
                expected = "þâÑÐŧ"
                assert expected.encode("latin-1") in content
                assert expected.encode("utf-8") in content
                gsos = content.split(b"strls")[1][1:-2]
                for gso in gsos.split(b"GSO")[1:]:
                    val = gso.split(b"\x00")[-2]
                    size = gso[gso.find(b"\x82") + 1]
                    assert len(val) == size - 1

    def test_encoding_latin1_118(self, datapath):
        # GH 25960
        msg = """
One or more strings in the dta file could not be decoded using utf-8, and
so the fallback encoding of latin-1 is being used.  This can happen when a file
has been incorrectly encoded by Stata or some other software. You should verify
the string values returned are correct."""
        # Move path outside of read_stata, or else assert_produces_warning
        # will block pytests skip mechanism from triggering (failing the test)
        # if the path is not present
        path = datapath("io", "data", "stata", "stata1_encoding_118.dta")
        with tm.assert_produces_warning(UnicodeWarning, filter_level="once") as w:
            encoded = read_stata(path)
            # with filter_level="always", produces 151 warnings which can be slow
            assert len(w) == 1
            assert w[0].message.args[0] == msg

        expected = DataFrame([["Düsseldorf"]] * 151, columns=["kreis1849"])
        tm.assert_frame_equal(encoded, expected)

    @pytest.mark.slow
    def test_stata_119(self, datapath):
        # Gzipped since contains 32,999 variables and uncompressed is 20MiB
        # Just validate that the reader reports correct number of variables
        # to avoid high peak memory
        with gzip.open(
            datapath("io", "data", "stata", "stata1_119.dta.gz"), "rb"
        ) as gz:
            with StataReader(gz) as reader:
                reader._ensure_open()
                assert reader._nvar == 32999

    @pytest.mark.parametrize("version", [118, 119, None])
    def test_utf8_writer(self, version):
        cat = pd.Categorical(["a", "β", "ĉ"], ordered=True)
        data = DataFrame(
            [
                [1.0, 1, "ᴬ", "ᴀ relatively long ŝtring"],
                [2.0, 2, "ᴮ", ""],
                [3.0, 3, "ᴰ", None],
            ],
            columns=["Å", "β", "ĉ", "strls"],
        )
        data["ᴐᴬᵀ"] = cat
        variable_labels = {
            "Å": "apple",
            "β": "ᵈᵉᵊ",
            "ĉ": "ᴎტჄႲႳႴႶႺ",
            "strls": "Long Strings",
            "ᴐᴬᵀ": "",
        }
        data_label = "ᴅaᵀa-label"
        value_labels = {"β": {1: "label", 2: "æøå", 3: "ŋot valid latin-1"}}
        data["β"] = data["β"].astype(np.int32)
        with tm.ensure_clean() as path:
            writer = StataWriterUTF8(
                path,
                data,
                data_label=data_label,
                convert_strl=["strls"],
                variable_labels=variable_labels,
                write_index=False,
                version=version,
                value_labels=value_labels,
            )
            writer.write_file()
            reread_encoded = read_stata(path)
            # Missing is intentionally converted to empty strl
            data["strls"] = data["strls"].fillna("")
            # Variable with value labels is reread as categorical
            data["β"] = (
                data["β"].replace(value_labels["β"]).astype("category").cat.as_ordered()
            )
            tm.assert_frame_equal(data, reread_encoded)
            with StataReader(path) as reader:
                assert reader.data_label == data_label
                assert reader.variable_labels() == variable_labels

            data.to_stata(path, version=version, write_index=False)
            reread_to_stata = read_stata(path)
            tm.assert_frame_equal(data, reread_to_stata)

    def test_writer_118_exceptions(self):
        df = DataFrame(np.zeros((1, 33000), dtype=np.int8))
        with tm.ensure_clean() as path:
            with pytest.raises(ValueError, match="version must be either 118 or 119."):
                StataWriterUTF8(path, df, version=117)
        with tm.ensure_clean() as path:
            with pytest.raises(ValueError, match="You must use version 119"):
                StataWriterUTF8(path, df, version=118)

    @pytest.mark.parametrize(
        "dtype_backend",
        ["numpy_nullable", pytest.param("pyarrow", marks=td.skip_if_no("pyarrow"))],
    )
    def test_read_write_ea_dtypes(self, dtype_backend):
        df = DataFrame(
            {
                "a": [1, 2, None],
                "b": ["a", "b", "c"],
                "c": [True, False, None],
                "d": [1.5, 2.5, 3.5],
                "e": pd.date_range("2020-12-31", periods=3, freq="D"),
            },
            index=pd.Index([0, 1, 2], name="index"),
        )
        df = df.convert_dtypes(dtype_backend=dtype_backend)
        df.to_stata("test_stata.dta", version=118)

        with tm.ensure_clean() as path:
            df.to_stata(path)
            written_and_read_again = self.read_dta(path)

        expected = DataFrame(
            {
                "a": [1, 2, np.nan],
                "b": ["a", "b", "c"],
                "c": [1.0, 0, np.nan],
                "d": [1.5, 2.5, 3.5],
                "e": pd.date_range("2020-12-31", periods=3, freq="D"),
            },
            index=pd.Index([0, 1, 2], name="index", dtype=np.int32),
        )

        tm.assert_frame_equal(written_and_read_again.set_index("index"), expected)


@pytest.mark.parametrize("version", [105, 108, 111, 113, 114])
def test_backward_compat(version, datapath):
    data_base = datapath("io", "data", "stata")
    ref = os.path.join(data_base, "stata-compat-118.dta")
    old = os.path.join(data_base, f"stata-compat-{version}.dta")
    expected = read_stata(ref)
    old_dta = read_stata(old)
    tm.assert_frame_equal(old_dta, expected, check_dtype=False)


def test_direct_read(datapath, monkeypatch):
    file_path = datapath("io", "data", "stata", "stata-compat-118.dta")

    # Test that opening a file path doesn't buffer the file.
    with StataReader(file_path) as reader:
        # Must not have been buffered to memory
        assert not reader.read().empty
        assert not isinstance(reader._path_or_buf, io.BytesIO)

    # Test that we use a given fp exactly, if possible.
    with open(file_path, "rb") as fp:
        with StataReader(fp) as reader:
            assert not reader.read().empty
            assert reader._path_or_buf is fp

    # Test that we use a given BytesIO exactly, if possible.
    with open(file_path, "rb") as fp:
        with io.BytesIO(fp.read()) as bio:
            with StataReader(bio) as reader:
                assert not reader.read().empty
                assert reader._path_or_buf is bio


def test_statareader_warns_when_used_without_context(datapath):
    file_path = datapath("io", "data", "stata", "stata-compat-118.dta")
    with tm.assert_produces_warning(
        ResourceWarning,
        match="without using a context manager",
    ):
        sr = StataReader(file_path)
        sr.read()
    with tm.assert_produces_warning(
        FutureWarning,
        match="is not part of the public API",
    ):
        sr.close()


@pytest.mark.parametrize("version", [114, 117, 118, 119, None])
@pytest.mark.parametrize("use_dict", [True, False])
@pytest.mark.parametrize("infer", [True, False])
def test_compression(compression, version, use_dict, infer, compression_to_extension):
    file_name = "dta_inferred_compression.dta"
    if compression:
        if use_dict:
            file_ext = compression
        else:
            file_ext = compression_to_extension[compression]
        file_name += f".{file_ext}"
    compression_arg = compression
    if infer:
        compression_arg = "infer"
    if use_dict:
        compression_arg = {"method": compression}

    df = DataFrame(
        np.random.default_rng(2).standard_normal((10, 2)), columns=list("AB")
    )
    df.index.name = "index"
    with tm.ensure_clean(file_name) as path:
        df.to_stata(path, version=version, compression=compression_arg)
        if compression == "gzip":
            with gzip.open(path, "rb") as comp:
                fp = io.BytesIO(comp.read())
        elif compression == "zip":
            with zipfile.ZipFile(path, "r") as comp:
                fp = io.BytesIO(comp.read(comp.filelist[0]))
        elif compression == "tar":
            with tarfile.open(path) as tar:
                fp = io.BytesIO(tar.extractfile(tar.getnames()[0]).read())
        elif compression == "bz2":
            with bz2.open(path, "rb") as comp:
                fp = io.BytesIO(comp.read())
        elif compression == "zstd":
            zstd = pytest.importorskip("zstandard")
            with zstd.open(path, "rb") as comp:
                fp = io.BytesIO(comp.read())
        elif compression == "xz":
            lzma = pytest.importorskip("lzma")
            with lzma.open(path, "rb") as comp:
                fp = io.BytesIO(comp.read())
        elif compression is None:
            fp = path
        reread = read_stata(fp, index_col="index")

    expected = df.copy()
    expected.index = expected.index.astype(np.int32)
    tm.assert_frame_equal(reread, expected)


@pytest.mark.parametrize("method", ["zip", "infer"])
@pytest.mark.parametrize("file_ext", [None, "dta", "zip"])
def test_compression_dict(method, file_ext):
    file_name = f"test.{file_ext}"
    archive_name = "test.dta"
    df = DataFrame(
        np.random.default_rng(2).standard_normal((10, 2)), columns=list("AB")
    )
    df.index.name = "index"
    with tm.ensure_clean(file_name) as path:
        compression = {"method": method, "archive_name": archive_name}
        df.to_stata(path, compression=compression)
        if method == "zip" or file_ext == "zip":
            with zipfile.ZipFile(path, "r") as zp:
                assert len(zp.filelist) == 1
                assert zp.filelist[0].filename == archive_name
                fp = io.BytesIO(zp.read(zp.filelist[0]))
        else:
            fp = path
        reread = read_stata(fp, index_col="index")

    expected = df.copy()
    expected.index = expected.index.astype(np.int32)
    tm.assert_frame_equal(reread, expected)


@pytest.mark.parametrize("version", [114, 117, 118, 119, None])
def test_chunked_categorical(version):
    df = DataFrame({"cats": Series(["a", "b", "a", "b", "c"], dtype="category")})
    df.index.name = "index"

    expected = df.copy()
    expected.index = expected.index.astype(np.int32)

    with tm.ensure_clean() as path:
        df.to_stata(path, version=version)
        with StataReader(path, chunksize=2, order_categoricals=False) as reader:
            for i, block in enumerate(reader):
                block = block.set_index("index")
                assert "cats" in block
                tm.assert_series_equal(
                    block.cats, expected.cats.iloc[2 * i : 2 * (i + 1)]
                )


def test_chunked_categorical_partial(datapath):
    dta_file = datapath("io", "data", "stata", "stata-dta-partially-labeled.dta")
    values = ["a", "b", "a", "b", 3.0]
    with StataReader(dta_file, chunksize=2) as reader:
        with tm.assert_produces_warning(CategoricalConversionWarning):
            for i, block in enumerate(reader):
                assert list(block.cats) == values[2 * i : 2 * (i + 1)]
                if i < 2:
                    idx = pd.Index(["a", "b"])
                else:
                    idx = pd.Index([3.0], dtype="float64")
                tm.assert_index_equal(block.cats.cat.categories, idx)
    with tm.assert_produces_warning(CategoricalConversionWarning):
        with StataReader(dta_file, chunksize=5) as reader:
            large_chunk = reader.__next__()
    direct = read_stata(dta_file)
    tm.assert_frame_equal(direct, large_chunk)


@pytest.mark.parametrize("chunksize", (-1, 0, "apple"))
def test_iterator_errors(datapath, chunksize):
    dta_file = datapath("io", "data", "stata", "stata-dta-partially-labeled.dta")
    with pytest.raises(ValueError, match="chunksize must be a positive"):
        with StataReader(dta_file, chunksize=chunksize):
            pass


def test_iterator_value_labels():
    # GH 31544
    values = ["c_label", "b_label"] + ["a_label"] * 500
    df = DataFrame({f"col{k}": pd.Categorical(values, ordered=True) for k in range(2)})
    with tm.ensure_clean() as path:
        df.to_stata(path, write_index=False)
        expected = pd.Index(["a_label", "b_label", "c_label"], dtype="object")
        with read_stata(path, chunksize=100) as reader:
            for j, chunk in enumerate(reader):
                for i in range(2):
                    tm.assert_index_equal(chunk.dtypes.iloc[i].categories, expected)
                tm.assert_frame_equal(chunk, df.iloc[j * 100 : (j + 1) * 100])


def test_precision_loss():
    df = DataFrame(
        [[sum(2**i for i in range(60)), sum(2**i for i in range(52))]],
        columns=["big", "little"],
    )
    with tm.ensure_clean() as path:
        with tm.assert_produces_warning(
            PossiblePrecisionLoss, match="Column converted from int64 to float64"
        ):
            df.to_stata(path, write_index=False)
        reread = read_stata(path)
        expected_dt = Series([np.float64, np.float64], index=["big", "little"])
        tm.assert_series_equal(reread.dtypes, expected_dt)
        assert reread.loc[0, "little"] == df.loc[0, "little"]
        assert reread.loc[0, "big"] == float(df.loc[0, "big"])


def test_compression_roundtrip(compression):
    df = DataFrame(
        [[0.123456, 0.234567, 0.567567], [12.32112, 123123.2, 321321.2]],
        index=["A", "B"],
        columns=["X", "Y", "Z"],
    )
    df.index.name = "index"

    with tm.ensure_clean() as path:
        df.to_stata(path, compression=compression)
        reread = read_stata(path, compression=compression, index_col="index")
        tm.assert_frame_equal(df, reread)

        # explicitly ensure file was compressed.
        with tm.decompress_file(path, compression) as fh:
            contents = io.BytesIO(fh.read())
        reread = read_stata(contents, index_col="index")
        tm.assert_frame_equal(df, reread)


@pytest.mark.parametrize("to_infer", [True, False])
@pytest.mark.parametrize("read_infer", [True, False])
def test_stata_compression(
    compression_only, read_infer, to_infer, compression_to_extension
):
    compression = compression_only

    ext = compression_to_extension[compression]
    filename = f"test.{ext}"

    df = DataFrame(
        [[0.123456, 0.234567, 0.567567], [12.32112, 123123.2, 321321.2]],
        index=["A", "B"],
        columns=["X", "Y", "Z"],
    )
    df.index.name = "index"

    to_compression = "infer" if to_infer else compression
    read_compression = "infer" if read_infer else compression

    with tm.ensure_clean(filename) as path:
        df.to_stata(path, compression=to_compression)
        result = read_stata(path, compression=read_compression, index_col="index")
        tm.assert_frame_equal(result, df)


def test_non_categorical_value_labels():
    data = DataFrame(
        {
            "fully_labelled": [1, 2, 3, 3, 1],
            "partially_labelled": [1.0, 2.0, np.nan, 9.0, np.nan],
            "Y": [7, 7, 9, 8, 10],
            "Z": pd.Categorical(["j", "k", "l", "k", "j"]),
        }
    )

    with tm.ensure_clean() as path:
        value_labels = {
            "fully_labelled": {1: "one", 2: "two", 3: "three"},
            "partially_labelled": {1.0: "one", 2.0: "two"},
        }
        expected = {**value_labels, "Z": {0: "j", 1: "k", 2: "l"}}

        writer = StataWriter(path, data, value_labels=value_labels)
        writer.write_file()

        with StataReader(path) as reader:
            reader_value_labels = reader.value_labels()
            assert reader_value_labels == expected

        msg = "Can't create value labels for notY, it wasn't found in the dataset."
        with pytest.raises(KeyError, match=msg):
            value_labels = {"notY": {7: "label1", 8: "label2"}}
            StataWriter(path, data, value_labels=value_labels)

        msg = (
            "Can't create value labels for Z, value labels "
            "can only be applied to numeric columns."
        )
        with pytest.raises(ValueError, match=msg):
            value_labels = {"Z": {1: "a", 2: "k", 3: "j", 4: "i"}}
            StataWriter(path, data, value_labels=value_labels)


def test_non_categorical_value_label_name_conversion():
    # Check conversion of invalid variable names
    data = DataFrame(
        {
            "invalid~!": [1, 1, 2, 3, 5, 8],  # Only alphanumeric and _
            "6_invalid": [1, 1, 2, 3, 5, 8],  # Must start with letter or _
            "invalid_name_longer_than_32_characters": [8, 8, 9, 9, 8, 8],  # Too long
            "aggregate": [2, 5, 5, 6, 6, 9],  # Reserved words
            (1, 2): [1, 2, 3, 4, 5, 6],  # Hashable non-string
        }
    )

    value_labels = {
        "invalid~!": {1: "label1", 2: "label2"},
        "6_invalid": {1: "label1", 2: "label2"},
        "invalid_name_longer_than_32_characters": {8: "eight", 9: "nine"},
        "aggregate": {5: "five"},
        (1, 2): {3: "three"},
    }

    expected = {
        "invalid__": {1: "label1", 2: "label2"},
        "_6_invalid": {1: "label1", 2: "label2"},
        "invalid_name_longer_than_32_char": {8: "eight", 9: "nine"},
        "_aggregate": {5: "five"},
        "_1__2_": {3: "three"},
    }

    with tm.ensure_clean() as path:
        with tm.assert_produces_warning(InvalidColumnName):
            data.to_stata(path, value_labels=value_labels)

        with StataReader(path) as reader:
            reader_value_labels = reader.value_labels()
            assert reader_value_labels == expected


def test_non_categorical_value_label_convert_categoricals_error():
    # Mapping more than one value to the same label is valid for Stata
    # labels, but can't be read with convert_categoricals=True
    value_labels = {
        "repeated_labels": {10: "Ten", 20: "More than ten", 40: "More than ten"}
    }

    data = DataFrame(
        {
            "repeated_labels": [10, 10, 20, 20, 40, 40],
        }
    )

    with tm.ensure_clean() as path:
        data.to_stata(path, value_labels=value_labels)

        with StataReader(path, convert_categoricals=False) as reader:
            reader_value_labels = reader.value_labels()
        assert reader_value_labels == value_labels

        col = "repeated_labels"
        repeats = "-" * 80 + "\n" + "\n".join(["More than ten"])

        msg = f"""
Value labels for column {col} are not unique. These cannot be converted to
pandas categoricals.

Either read the file with `convert_categoricals` set to False or use the
low level interface in `StataReader` to separately read the values and the
value_labels.

The repeated labels are:
{repeats}
"""
        with pytest.raises(ValueError, match=msg):
            read_stata(path, convert_categoricals=True)


@pytest.mark.parametrize("version", [114, 117, 118, 119, None])
@pytest.mark.parametrize(
    "dtype",
    [
        pd.BooleanDtype,
        pd.Int8Dtype,
        pd.Int16Dtype,
        pd.Int32Dtype,
        pd.Int64Dtype,
        pd.UInt8Dtype,
        pd.UInt16Dtype,
        pd.UInt32Dtype,
        pd.UInt64Dtype,
    ],
)
def test_nullable_support(dtype, version):
    df = DataFrame(
        {
            "a": Series([1.0, 2.0, 3.0]),
            "b": Series([1, pd.NA, pd.NA], dtype=dtype.name),
            "c": Series(["a", "b", None]),
        }
    )
    dtype_name = df.b.dtype.numpy_dtype.name
    # Only use supported names: no uint, bool or int64
    dtype_name = dtype_name.replace("u", "")
    if dtype_name == "int64":
        dtype_name = "int32"
    elif dtype_name == "bool":
        dtype_name = "int8"
    value = StataMissingValue.BASE_MISSING_VALUES[dtype_name]
    smv = StataMissingValue(value)
    expected_b = Series([1, smv, smv], dtype=object, name="b")
    expected_c = Series(["a", "b", ""], name="c")
    with tm.ensure_clean() as path:
        df.to_stata(path, write_index=False, version=version)
        reread = read_stata(path, convert_missing=True)
        tm.assert_series_equal(df.a, reread.a)
        tm.assert_series_equal(reread.b, expected_b)
        tm.assert_series_equal(reread.c, expected_c)


def test_empty_frame():
    # GH 46240
    # create an empty DataFrame with int64 and float64 dtypes
    df = DataFrame(data={"a": range(3), "b": [1.0, 2.0, 3.0]}).head(0)
    with tm.ensure_clean() as path:
        df.to_stata(path, write_index=False, version=117)
        # Read entire dataframe
        df2 = read_stata(path)
        assert "b" in df2
        # Dtypes don't match since no support for int32
        dtypes = Series({"a": np.dtype("int32"), "b": np.dtype("float64")})
        tm.assert_series_equal(df2.dtypes, dtypes)
        # read one column of empty .dta file
        df3 = read_stata(path, columns=["a"])
        assert "b" not in df3
        tm.assert_series_equal(df3.dtypes, dtypes.loc[["a"]])