2382 lines
90 KiB
Python
2382 lines
90 KiB
Python
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 __<5F>__ 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"]])
|