Coverage for /pythoncovmergedfiles/medio/medio/usr/local/lib/python3.11/site-packages/pandas/io/stata.py: 14%
Shortcuts on this page
r m x toggle line displays
j k next/prev highlighted chunk
0 (zero) top of page
1 (one) first highlighted chunk
1"""
2Module contains tools for processing Stata files into DataFrames
4The StataReader below was originally written by Joe Presbrey as part of PyDTA.
5It has been extended and improved by Skipper Seabold from the Statsmodels
6project who also developed the StataWriter and was finally added to pandas in
7a once again improved version.
9You can find more information on http://presbrey.mit.edu/PyDTA and
10https://www.statsmodels.org/devel/
11"""
13from __future__ import annotations
15from collections import abc
16from datetime import (
17 datetime,
18 timedelta,
19)
20from io import BytesIO
21import os
22import struct
23import sys
24from typing import (
25 IO,
26 TYPE_CHECKING,
27 AnyStr,
28 Final,
29 Self,
30)
31import warnings
33import numpy as np
35from pandas._libs import lib
36from pandas._libs.lib import infer_dtype
37from pandas._libs.writers import max_len_string_array
38from pandas.errors import (
39 CategoricalConversionWarning,
40 InvalidColumnName,
41 OutOfBoundsDatetime,
42 Pandas4Warning,
43 PossiblePrecisionLoss,
44 ValueLabelTypeMismatch,
45)
46from pandas.util._decorators import (
47 set_module,
48)
49from pandas.util._exceptions import find_stack_level
51from pandas.core.dtypes.base import ExtensionDtype
52from pandas.core.dtypes.common import (
53 ensure_object,
54 is_numeric_dtype,
55 is_string_dtype,
56)
57from pandas.core.dtypes.dtypes import CategoricalDtype
59from pandas import (
60 Categorical,
61 DatetimeIndex,
62 NaT,
63 Timestamp,
64 isna,
65 to_datetime,
66)
67from pandas.core.frame import DataFrame
68from pandas.core.indexes.base import Index
69from pandas.core.indexes.range import RangeIndex
70from pandas.core.series import Series
72from pandas.io.common import get_handle
74if TYPE_CHECKING:
75 from collections.abc import (
76 Callable,
77 Hashable,
78 Sequence,
79 )
80 from types import TracebackType
81 from typing import Literal
83 from pandas._typing import (
84 CompressionOptions,
85 FilePath,
86 ReadBuffer,
87 StorageOptions,
88 WriteBuffer,
89 )
91# The format versions each header reader accepts. _version_error enumerates
92# both, so keep them in sync -- GH#63082 dropped four of these from the
93# message while the readers went on accepting them.
94_old_format_versions: Final = [102, 103, 104, 105, 108, 110, 111, 113, 114, 115]
95_new_format_versions: Final = [117, 118, 119]
97# Error shown when a version number was parsed but is not supported.
98# Wording intentionally mentions “either not a valid Stata dataset or
99# an unsupported version” to avoid confusing users when input is not a
100# real .dta.
101_version_error = (
102 "This is either not a valid Stata dataset or a Stata dataset from a "
103 "version pandas does not support (detected: {version}). pandas "
104 "supports importing versions 102, 103, 104, 105, 108, 110 (Stata 7), "
105 "111 (Stata 7SE), 113 (Stata 8/9), 114 (Stata 10/11), 115 (Stata 12), "
106 "117 (Stata 13), 118 (Stata 14/15/16), and 119 (Stata 15/16, over "
107 "32,767 variables)."
108)
111_date_formats = ["%tc", "%tC", "%td", "%d", "%tw", "%tm", "%tq", "%th", "%ty"]
114stata_epoch: Final = datetime(1960, 1, 1)
115unix_epoch: Final = datetime(1970, 1, 1)
118def _stata_ordinals_to_datetime64(
119 ordinals: Series | np.ndarray,
120 fmt: str,
121 unit: str,
122 shift: np.timedelta64 | None = None,
123 mult: int = 1,
124) -> np.ndarray:
125 """
126 Cast Stata date ordinals to ``datetime64[unit]``, scaling by `mult` and
127 then adding `shift`.
129 Casting a float to datetime64 saturates at the int64 bounds instead of
130 raising, so an ordinal too far from the epoch would otherwise decode to a
131 bogus date. Reject those up front, naming the caller's own ordinal rather
132 than the scaled one, and let ``Series`` do the range-checked conversion to
133 the caller's resolution (GH#36096).
134 """
135 values = np.asarray(ordinals)
136 if values.dtype.kind == "f":
137 # The scaled ordinal has to survive the int64 cast on its own, and the
138 # shifted result has to fit too; `shift` is one-signed, so it buys
139 # headroom at one end and costs it at the other. Take both bounds.
140 # int64 min doubles as the NaT sentinel, so that end is exclusive.
141 offset = 0 if shift is None else shift.astype("i8")
142 lower = max(-(2.0**63), -(2.0**63) - offset) / mult
143 upper = min(2.0**63, 2.0**63 - offset) / mult
144 # NaN compares False both ways and casts to NaT, which is what we want.
145 out_of_range = (values <= lower) | (values >= upper)
146 if out_of_range.any():
147 raise OutOfBoundsDatetime(
148 f"Out of bounds date value {values[out_of_range][0]} for format {fmt}"
149 )
150 if mult != 1:
151 values = values * mult
152 res = values.astype(f"M8[{unit}]")
153 if shift is not None:
154 res += shift
155 return res
158def _stata_elapsed_date_to_datetime_vec(dates: Series, fmt: str) -> Series:
159 """
160 Convert from SIF to datetime. https://www.stata.com/help.cgi?datetime
162 Parameters
163 ----------
164 dates : Series
165 The Stata Internal Format date to convert to datetime according to fmt
166 fmt : str
167 The format to convert to. Can be, tc, td, tw, tm, tq, th, ty
168 Returns
170 Returns
171 -------
172 converted : Series
173 The converted dates
175 Examples
176 --------
177 >>> dates = pd.Series([52])
178 >>> _stata_elapsed_date_to_datetime_vec(dates, "%tw")
179 0 1961-01-01
180 dtype: datetime64[s]
182 Notes
183 -----
184 datetime/c - tc
185 milliseconds since 01jan1960 00:00:00.000, assuming 86,400 s/day
186 datetime/C - tC - NOT IMPLEMENTED
187 milliseconds since 01jan1960 00:00:00.000, adjusted for leap seconds
188 date - td
189 days since 01jan1960 (01jan1960 = 0)
190 weekly date - tw
191 weeks since 1960w1
192 This assumes 52 weeks in a year, then adds 7 * remainder of the weeks.
193 The datetime value is the start of the week in terms of days in the
194 year, not ISO calendar weeks.
195 monthly date - tm
196 months since 1960m1
197 quarterly date - tq
198 quarters since 1960q1
199 half-yearly date - th
200 half-years since 1960h1 yearly
201 date - ty
202 years since 0000
203 """
205 if fmt.startswith(("%tc", "tc")):
206 # Delta ms relative to base
207 td = np.timedelta64(stata_epoch - unix_epoch, "ms")
208 res = _stata_ordinals_to_datetime64(dates._values, fmt, "ms", td)
209 return Series(res, index=dates.index)
211 elif fmt.startswith(("%td", "td", "%d", "d")):
212 # Delta days relative to base
213 td = np.timedelta64(stata_epoch - unix_epoch, "D")
214 res = _stata_ordinals_to_datetime64(dates._values, fmt, "D", td)
215 return Series(res, index=dates.index)
217 elif fmt.startswith(("%tm", "tm")):
218 # Delta months relative to base. Stata stores these in anything from a
219 # byte to a double, so widen before shifting the epoch: the narrow
220 # integer types overflow on the way (GH#36096).
221 ordinals = dates.astype("f8") + (stata_epoch.year - unix_epoch.year) * 12
222 res = _stata_ordinals_to_datetime64(ordinals, fmt, "M")
223 return Series(res, index=dates.index)
225 elif fmt.startswith(("%tq", "tq")):
226 # Delta quarters relative to base, expressed in months because the
227 # conversion below reads datetime64[3M] as datetime64[M]. Truncate to
228 # whole quarters first, as the datetime64[3M] cast used to.
229 quarters = dates.astype("f8") + (stata_epoch.year - unix_epoch.year) * 4
230 res = _stata_ordinals_to_datetime64(np.trunc(quarters), fmt, "M", mult=3)
231 return Series(res, index=dates.index)
233 elif fmt.startswith(("%th", "th")):
234 # Delta half-years relative to base, expressed in months (see %tq)
235 halves = dates.astype("f8") + (stata_epoch.year - unix_epoch.year) * 2
236 res = _stata_ordinals_to_datetime64(np.trunc(halves), fmt, "M", mult=6)
237 return Series(res, index=dates.index)
239 elif fmt.startswith(("%ty", "ty")):
240 # Years -- not delta
241 ordinals = dates.astype("f8") - 1970
242 res = _stata_ordinals_to_datetime64(ordinals, fmt, "Y")
243 return Series(res, index=dates.index)
245 if fmt.startswith(("%tC", "tC")):
246 warnings.warn(
247 "Encountered %tC format. Leaving in Stata Internal Format.",
248 stacklevel=find_stack_level(),
249 )
250 # Converts nothing, so hand back the file's own numbers; this must stay
251 # above the int64 cast below, which saturates out-of-bounds values
252 # (GH#68034).
253 conv_dates = Series(dates, dtype=object)
254 na_locs = isna(dates)
255 if na_locs.any():
256 conv_dates[na_locs] = NaT
257 return conv_dates
259 bad_locs = np.isnan(dates)
260 has_bad_values = False
261 if bad_locs.any():
262 has_bad_values = True
263 dates._values[bad_locs] = 1.0 # Replace with NaT
264 dates = dates.astype(np.int64)
266 # does not count leap days - 7 days is a week.
267 # 52nd week may have more than 7 days
268 if fmt.startswith(("%tw", "tw")):
269 year = stata_epoch.year + dates // 52
270 days = (dates % 52) * 7
271 res = _stata_ordinals_to_datetime64(year - 1970, fmt, "Y")
272 conv_dates = Series(res, index=dates.index) + days._values.astype("m8[D]")
274 else:
275 raise ValueError(f"Date fmt {fmt} not understood")
277 if has_bad_values: # Restore NaT for bad values
278 conv_dates[bad_locs] = NaT
280 return conv_dates
283def _datetime_to_stata_elapsed_vec(dates: Series, fmt: str) -> Series:
284 """
285 Convert from datetime to SIF. https://www.stata.com/help.cgi?datetime
287 Parameters
288 ----------
289 dates : Series
290 Series or array containing datetime or datetime64[ns] to
291 convert to the Stata Internal Format given by fmt
292 fmt : str
293 The format to convert to. Can be, tc, td, tw, tm, tq, th, ty
294 """
295 index = dates.index
296 NS_PER_DAY = 24 * 3600 * 1000 * 1000 * 1000
297 US_PER_DAY = NS_PER_DAY / 1000
298 MS_PER_DAY = NS_PER_DAY / 1_000_000
300 def parse_dates_safe(
301 dates: Series, delta: bool = False, year: bool = False, days: bool = False
302 ) -> DataFrame:
303 d = {}
304 if lib.is_np_dtype(dates.dtype, "M"):
305 if delta:
306 time_delta = dates.dt.as_unit("ms") - Timestamp(stata_epoch).as_unit(
307 "ms"
308 )
309 d["delta"] = time_delta._values.view(np.int64)
310 if days or year:
311 date_index = DatetimeIndex(dates)
312 d["year"] = date_index._data.year
313 d["month"] = date_index._data.month
314 if days:
315 year_start = np.asarray(dates).astype("M8[Y]").astype(dates.dtype)
316 diff = dates - year_start
317 d["days"] = np.asarray(diff).astype("m8[D]").view("int64")
319 elif infer_dtype(dates, skipna=False) == "datetime":
320 warnings.warn(
321 # GH#56536
322 "Converting object-dtype columns of datetimes to datetime64 when "
323 "writing to stata is deprecated. Call "
324 "`df=df.infer_objects(copy=False)` before writing to stata instead.",
325 Pandas4Warning,
326 stacklevel=find_stack_level(),
327 )
328 if delta:
329 delta = dates._values - stata_epoch
331 def f(x: timedelta) -> float:
332 return US_PER_DAY * x.days + 1_000_000 * x.seconds + x.microseconds
334 v = np.vectorize(f)
335 d["delta"] = v(delta) // 1_000 # convert back to ms
336 if year:
337 try:
338 date_index = DatetimeIndex(dates)
339 except ValueError:
340 # GH#64556 entries not sharing a single tzinfo cannot be
341 # cast to datetime64, but we only need each object's fields
342 d["year"] = [date.year for date in dates]
343 d["month"] = [date.month for date in dates]
344 else:
345 d["year"] = date_index.year
346 d["month"] = date_index.month
347 if days:
349 def g(x: datetime) -> int:
350 return (x - datetime(x.year, 1, 1)).days
352 v = np.vectorize(g)
353 d["days"] = v(dates)
354 else:
355 raise ValueError(
356 "Columns containing dates must contain either "
357 "datetime64, datetime or null values."
358 )
360 return DataFrame(d, index=index)
362 bad_loc = isna(dates)
363 index = dates.index
364 if bad_loc.any():
365 if lib.is_np_dtype(dates.dtype, "M"):
366 dates._values[bad_loc] = to_datetime(stata_epoch)
367 else:
368 dates._values[bad_loc] = stata_epoch
370 if fmt in ["%tc", "tc"]:
371 d = parse_dates_safe(dates, delta=True)
372 conv_dates = d.delta
373 elif fmt in ["%tC", "tC"]:
374 warnings.warn(
375 "Stata Internal Format tC not supported.",
376 stacklevel=find_stack_level(),
377 )
378 conv_dates = dates
379 elif fmt in ["%td", "td"]:
380 d = parse_dates_safe(dates, delta=True)
381 conv_dates = d.delta // MS_PER_DAY
382 elif fmt in ["%tw", "tw"]:
383 d = parse_dates_safe(dates, year=True, days=True)
384 conv_dates = 52 * (d.year - stata_epoch.year) + d.days // 7
385 elif fmt in ["%tm", "tm"]:
386 d = parse_dates_safe(dates, year=True)
387 conv_dates = 12 * (d.year - stata_epoch.year) + d.month - 1
388 elif fmt in ["%tq", "tq"]:
389 d = parse_dates_safe(dates, year=True)
390 conv_dates = 4 * (d.year - stata_epoch.year) + (d.month - 1) // 3
391 elif fmt in ["%th", "th"]:
392 d = parse_dates_safe(dates, year=True)
393 conv_dates = 2 * (d.year - stata_epoch.year) + (d.month > 6).astype(int)
394 elif fmt in ["%ty", "ty"]:
395 d = parse_dates_safe(dates, year=True)
396 conv_dates = d.year
397 else:
398 raise ValueError(f"Format {fmt} is not a known Stata date format")
400 conv_dates = Series(conv_dates, dtype=np.float64, copy=False)
401 missing_value = struct.unpack("<d", b"\x00\x00\x00\x00\x00\x00\xe0\x7f")[0]
402 conv_dates[bad_loc] = missing_value
404 return Series(conv_dates, index=index, copy=False)
407excessive_string_length_error: Final = """
408Fixed width strings in Stata .dta files are limited to 244 (or fewer)
409characters. Column '{0}' does not satisfy this restriction. Use the
410'version=117' parameter to write the newer (Stata 13 and later) format.
411"""
414precision_loss_doc: Final = """
415Column converted from {0} to {1}, and some data are outside of the lossless
416conversion range. This may result in a loss of precision in the saved data.
417"""
420value_label_mismatch_doc: Final = """
421Stata value labels (pandas categories) must be strings. Column {0} contains
422non-string labels which will be converted to strings. Please check that the
423Stata data file created has not lost information due to duplicate labels.
424"""
427invalid_name_doc: Final = """
428Not all pandas column names were valid Stata variable names.
429The following replacements have been made:
431 {0}
433If this is not what you expect, please make sure you have Stata-compliant
434column names in your DataFrame (strings only, max 32 characters, only
435alphanumerics and underscores, no Stata reserved words)
436"""
439categorical_conversion_warning: Final = """
440One or more series with value labels are not fully labeled. Reading this
441dataset with an iterator results in categorical variable with different
442categories. This occurs since it is not possible to know all possible values
443until the entire dataset has been read. To avoid this warning, you can either
444read dataset without an iterator, or manually convert categorical data by
445``convert_categoricals`` to False and then accessing the variable labels
446through the value_labels method of the reader.
447"""
450def _cast_to_stata_types(data: DataFrame) -> DataFrame:
451 """
452 Checks the dtypes of the columns of a pandas DataFrame for
453 compatibility with the data types and ranges supported by Stata, and
454 converts if necessary.
456 Parameters
457 ----------
458 data : DataFrame
459 The DataFrame to check and convert
461 Notes
462 -----
463 Numeric columns in Stata must be one of int8, int16, int32, float32 or
464 float64, with some additional value restrictions. int8 and int16 columns
465 are checked for violations of the value restrictions and upcast if needed.
466 int64 data is not usable in Stata, and so it is downcast to int32 whenever
467 the value are in the int32 range, and sidecast to float64 when larger than
468 this range. If the int64 values are outside of the range of those
469 perfectly representable as float64 values, a warning is raised.
471 bool columns are cast to int8. uint columns are converted to int of the
472 same size if there is no loss in precision, otherwise are upcast to a
473 larger type. uint64 is currently not supported since it is concerted to
474 object in a DataFrame.
475 """
476 ws = ""
477 # original, if small, if large
478 conversion_data: tuple[
479 tuple[type, type, type],
480 tuple[type, type, type],
481 tuple[type, type, type],
482 tuple[type, type, type],
483 tuple[type, type, type],
484 ] = (
485 (np.bool_, np.int8, np.int8),
486 (np.uint8, np.int8, np.int16),
487 (np.uint16, np.int16, np.int32),
488 (np.uint32, np.int32, np.int64),
489 (np.uint64, np.int64, np.float64),
490 )
492 float32_max = struct.unpack("<f", b"\xff\xff\xff\x7e")[0]
493 float64_max = struct.unpack("<d", b"\xff\xff\xff\xff\xff\xff\xdf\x7f")[0]
495 for col in data:
496 # Cast from unsupported types to supported types
497 is_nullable_int = (
498 isinstance(data[col].dtype, ExtensionDtype)
499 and data[col].dtype.kind in "iub"
500 )
501 # We need to find orig_missing before altering data below
502 orig_missing = data[col].isna()
503 if is_nullable_int:
504 fv = 0 if data[col].dtype.kind in "iu" else False
505 # Replace with NumPy-compatible column
506 data[col] = data[col].fillna(fv).astype(data[col].dtype.numpy_dtype)
507 elif isinstance(data[col].dtype, ExtensionDtype):
508 if getattr(data[col].dtype, "numpy_dtype", None) is not None:
509 data[col] = data[col].astype(data[col].dtype.numpy_dtype)
510 elif is_string_dtype(data[col].dtype):
511 # TODO could avoid converting string dtype to object here,
512 # but handle string dtype in _encode_strings
513 data[col] = data[col].astype("object")
514 # generate_table checks for None values
515 data.loc[data[col].isna(), col] = None
517 dtype = data[col].dtype
518 empty_df = data.shape[0] == 0
519 for c_data in conversion_data:
520 if dtype == c_data[0]:
521 if empty_df or data[col].max() <= np.iinfo(c_data[1]).max:
522 dtype = c_data[1]
523 else:
524 dtype = c_data[2]
525 if c_data[2] == np.int64: # Warn if necessary
526 if data[col].max() >= 2**53:
527 ws = precision_loss_doc.format("uint64", "float64")
529 data[col] = data[col].astype(dtype)
531 # Check values and upcast if necessary
533 if dtype == np.int8 and not empty_df:
534 if data[col].max() > 100 or data[col].min() < -127:
535 data[col] = data[col].astype(np.int16)
536 elif dtype == np.int16 and not empty_df:
537 if data[col].max() > 32740 or data[col].min() < -32767:
538 data[col] = data[col].astype(np.int32)
539 elif dtype == np.int64:
540 if empty_df or (
541 data[col].max() <= 2147483620 and data[col].min() >= -2147483647
542 ):
543 data[col] = data[col].astype(np.int32)
544 else:
545 data[col] = data[col].astype(np.float64)
546 if data[col].max() >= 2**53 or data[col].min() <= -(2**53):
547 ws = precision_loss_doc.format("int64", "float64")
548 elif dtype in (np.float32, np.float64):
549 if np.isinf(data[col]).any():
550 raise ValueError(
551 f"Column {col} contains infinity or -infinity"
552 "which is outside the range supported by Stata."
553 )
554 value = data[col].max()
555 if dtype == np.float32 and value > float32_max:
556 data[col] = data[col].astype(np.float64)
557 elif dtype == np.float64:
558 if value > float64_max:
559 raise ValueError(
560 f"Column {col} has a maximum value ({value}) outside the range "
561 f"supported by Stata ({float64_max})"
562 )
563 if is_nullable_int:
564 if orig_missing.any():
565 # Replace missing by Stata sentinel value
566 sentinel = StataMissingValue.BASE_MISSING_VALUES[data[col].dtype.name]
567 data.loc[orig_missing, col] = sentinel
568 if ws:
569 warnings.warn(
570 ws,
571 PossiblePrecisionLoss,
572 stacklevel=find_stack_level(),
573 )
575 return data
578class StataValueLabel:
579 """
580 Parse a categorical column and prepare formatted output
582 Parameters
583 ----------
584 catarray : Series
585 Categorical Series to encode
586 encoding : {"latin-1", "utf-8"}
587 Encoding to use for value labels.
588 """
590 def __init__(
591 self, catarray: Series, encoding: Literal["latin-1", "utf-8"] = "latin-1"
592 ) -> None:
593 if encoding not in ("latin-1", "utf-8"):
594 raise ValueError("Only latin-1 and utf-8 are supported.")
595 self.labname = catarray.name
596 self._encoding = encoding
597 categories = catarray.cat.categories
598 self.value_labels = enumerate(categories)
600 self._prepare_value_labels()
602 def _prepare_value_labels(self) -> None:
603 """Encode value labels."""
605 self.text_len = 0
606 self.txt: list[bytes] = []
607 self.n = 0
608 # Offsets (length of categories), converted to int32
609 self.off = np.array([], dtype=np.int32)
610 # Values, converted to int32
611 self.val = np.array([], dtype=np.int32)
612 self.len = 0
614 # Compute lengths and setup lists of offsets and labels
615 offsets: list[int] = []
616 values: list[float] = []
617 for vl in self.value_labels:
618 category: str | bytes = vl[1]
619 if not isinstance(category, str):
620 category = str(category)
621 warnings.warn(
622 value_label_mismatch_doc.format(self.labname),
623 ValueLabelTypeMismatch,
624 stacklevel=find_stack_level(),
625 )
626 category = category.encode(self._encoding)
627 offsets.append(self.text_len)
628 self.text_len += len(category) + 1 # +1 for the padding
629 values.append(vl[0])
630 self.txt.append(category)
631 self.n += 1
633 # Ensure int32
634 self.off = np.array(offsets, dtype=np.int32)
635 self.val = np.array(values, dtype=np.int32)
637 # Total length
638 self.len = 4 + 4 + 4 * self.n + 4 * self.n + self.text_len
640 def generate_value_label(self, byteorder: str) -> bytes:
641 """
642 Generate the binary representation of the value labels.
644 Parameters
645 ----------
646 byteorder : str
647 Byte order of the output
649 Returns
650 -------
651 value_label : bytes
652 Bytes containing the formatted value label
653 """
654 encoding = self._encoding
655 bio = BytesIO()
656 null_byte = b"\x00"
658 # len
659 bio.write(struct.pack(byteorder + "i", self.len))
661 # labname
662 labname = str(self.labname)[:32].encode(encoding)
663 lab_len = 32 if encoding not in ("utf-8", "utf8") else 128
664 labname = _pad_bytes(labname, lab_len + 1)
665 bio.write(labname)
667 # padding - 3 bytes
668 for i in range(3):
669 bio.write(struct.pack("c", null_byte))
671 # value_label_table
672 # n - int32
673 bio.write(struct.pack(byteorder + "i", self.n))
675 # textlen - int32
676 bio.write(struct.pack(byteorder + "i", self.text_len))
678 # off - int32 array (n elements)
679 for offset in self.off:
680 bio.write(struct.pack(byteorder + "i", offset))
682 # val - int32 array (n elements)
683 for value in self.val:
684 bio.write(struct.pack(byteorder + "i", value))
686 # txt - Text labels, null terminated
687 for text in self.txt:
688 bio.write(text + null_byte)
690 return bio.getvalue()
693class StataNonCatValueLabel(StataValueLabel):
694 """
695 Prepare formatted version of value labels
697 Parameters
698 ----------
699 labname : str
700 Value label name
701 value_labels: Dictionary
702 Mapping of values to labels
703 encoding : {"latin-1", "utf-8"}
704 Encoding to use for value labels.
705 """
707 def __init__(
708 self,
709 labname: str,
710 value_labels: dict[float, str],
711 encoding: Literal["latin-1", "utf-8"] = "latin-1",
712 ) -> None:
713 if encoding not in ("latin-1", "utf-8"):
714 raise ValueError("Only latin-1 and utf-8 are supported.")
716 self.labname = labname
717 self._encoding = encoding
718 self.value_labels = sorted( # type: ignore[assignment]
719 value_labels.items(), key=lambda x: x[0]
720 )
721 self._prepare_value_labels()
724class StataMissingValue:
725 """
726 An observation's missing value.
728 Parameters
729 ----------
730 value : {int, float}
731 The Stata missing value code
733 Notes
734 -----
735 More information: <https://www.stata.com/help.cgi?missing>
737 Integer missing values make the code '.', '.a', ..., '.z' to the ranges
738 101 ... 127 (for int8), 32741 ... 32767 (for int16) and 2147483621 ...
739 2147483647 (for int32). Missing values for floating point data types are
740 more complex but the pattern is simple to discern from the following table.
742 np.float32 missing values (float in Stata)
743 0000007f .
744 0008007f .a
745 0010007f .b
746 ...
747 00c0007f .x
748 00c8007f .y
749 00d0007f .z
751 np.float64 missing values (double in Stata)
752 000000000000e07f .
753 000000000001e07f .a
754 000000000002e07f .b
755 ...
756 000000000018e07f .x
757 000000000019e07f .y
758 00000000001ae07f .z
759 """
761 # Construct a dictionary of missing values
762 MISSING_VALUES: dict[float, str] = {}
763 bases: Final = (101, 32741, 2147483621)
764 for b in bases:
765 # Conversion to long to avoid hash issues on 32 bit platforms #8968
766 MISSING_VALUES[b] = "."
767 for i in range(1, 27):
768 MISSING_VALUES[i + b] = "." + chr(96 + i)
770 float32_base: bytes = b"\x00\x00\x00\x7f"
771 increment_32: int = struct.unpack("<i", b"\x00\x08\x00\x00")[0]
772 for i in range(27):
773 key = struct.unpack("<f", float32_base)[0]
774 MISSING_VALUES[key] = "."
775 if i > 0:
776 MISSING_VALUES[key] += chr(96 + i)
777 int_value = struct.unpack("<i", struct.pack("<f", key))[0] + increment_32
778 float32_base = struct.pack("<i", int_value)
780 float64_base: bytes = b"\x00\x00\x00\x00\x00\x00\xe0\x7f"
781 increment_64 = struct.unpack("q", b"\x00\x00\x00\x00\x00\x01\x00\x00")[0]
782 for i in range(27):
783 key = struct.unpack("<d", float64_base)[0]
784 MISSING_VALUES[key] = "."
785 if i > 0:
786 MISSING_VALUES[key] += chr(96 + i)
787 int_value = struct.unpack("q", struct.pack("<d", key))[0] + increment_64
788 float64_base = struct.pack("q", int_value)
790 BASE_MISSING_VALUES: Final = {
791 "int8": 101,
792 "int16": 32741,
793 "int32": 2147483621,
794 "float32": struct.unpack("<f", float32_base)[0],
795 "float64": struct.unpack("<d", float64_base)[0],
796 }
798 def __init__(self, value: float) -> None:
799 self._value = value
800 # Conversion to int to avoid hash issues on 32 bit platforms #8968
801 value = int(value) if value < 2147483648 else float(value)
802 self._str = self.MISSING_VALUES[value]
804 @property
805 def string(self) -> str:
806 """
807 The Stata representation of the missing value: '.', '.a'..'.z'
809 Returns
810 -------
811 str
812 The representation of the missing value.
813 """
814 return self._str
816 @property
817 def value(self) -> float:
818 """
819 The binary representation of the missing value.
821 Returns
822 -------
823 {int, float}
824 The binary representation of the missing value.
825 """
826 return self._value
828 def __str__(self) -> str:
829 return self.string
831 def __repr__(self) -> str:
832 return f"{type(self)}({self})"
834 def __eq__(self, other: object) -> bool:
835 return (
836 isinstance(other, type(self))
837 and self.string == other.string
838 and self.value == other.value
839 )
841 @classmethod
842 def get_base_missing_value(cls, dtype: np.dtype) -> float:
843 if dtype.type is np.int8:
844 value = cls.BASE_MISSING_VALUES["int8"]
845 elif dtype.type is np.int16:
846 value = cls.BASE_MISSING_VALUES["int16"]
847 elif dtype.type is np.int32:
848 value = cls.BASE_MISSING_VALUES["int32"]
849 elif dtype.type is np.float32:
850 value = cls.BASE_MISSING_VALUES["float32"]
851 elif dtype.type is np.float64:
852 value = cls.BASE_MISSING_VALUES["float64"]
853 else:
854 raise ValueError("Unsupported dtype")
855 return value
858class StataParser:
859 def __init__(self) -> None:
860 # type code.
861 # --------------------
862 # str1 1 = 0x01
863 # str2 2 = 0x02
864 # ...
865 # str244 244 = 0xf4
866 # byte 251 = 0xfb (sic)
867 # int 252 = 0xfc
868 # long 253 = 0xfd
869 # float 254 = 0xfe
870 # double 255 = 0xff
871 # --------------------
872 # NOTE: the byte type seems to be reserved for categorical variables
873 # with a label, but the underlying variable is -127 to 100
874 # we're going to drop the label and cast to int
875 self.DTYPE_MAP = dict(
876 [(i, np.dtype(f"S{i}")) for i in range(1, 245)]
877 + [
878 (251, np.dtype(np.int8)),
879 (252, np.dtype(np.int16)),
880 (253, np.dtype(np.int32)),
881 (254, np.dtype(np.float32)),
882 (255, np.dtype(np.float64)),
883 ]
884 )
885 self.DTYPE_MAP_XML: dict[int, np.dtype] = {
886 32768: np.dtype(np.uint8), # Keys to GSO
887 65526: np.dtype(np.float64),
888 65527: np.dtype(np.float32),
889 65528: np.dtype(np.int32),
890 65529: np.dtype(np.int16),
891 65530: np.dtype(np.int8),
892 }
893 self.TYPE_MAP = list(tuple(range(251)) + tuple("bhlfd"))
894 self.TYPE_MAP_XML = {
895 # Not really a Q, unclear how to handle byteswap
896 32768: "Q",
897 65526: "d",
898 65527: "f",
899 65528: "l",
900 65529: "h",
901 65530: "b",
902 }
903 # NOTE: technically, some of these are wrong. there are more numbers
904 # that can be represented. it's the 27 ABOVE and BELOW the max listed
905 # numeric data type in [U] 12.2.2 of the 11.2 manual
906 float32_min = b"\xff\xff\xff\xfe"
907 float32_max = b"\xff\xff\xff\x7e"
908 float64_min = b"\xff\xff\xff\xff\xff\xff\xef\xff"
909 float64_max = b"\xff\xff\xff\xff\xff\xff\xdf\x7f"
910 self.VALID_RANGE = {
911 "b": (-127, 100),
912 "h": (-32767, 32740),
913 "l": (-2147483647, 2147483620),
914 "f": (
915 np.float32(struct.unpack("<f", float32_min)[0]),
916 np.float32(struct.unpack("<f", float32_max)[0]),
917 ),
918 "d": (
919 np.float64(struct.unpack("<d", float64_min)[0]),
920 np.float64(struct.unpack("<d", float64_max)[0]),
921 ),
922 }
923 self.OLD_VALID_RANGE = {
924 "b": (-128, 126),
925 "h": (-32768, 32766),
926 "l": (-2147483648, 2147483646),
927 "f": (
928 np.float32(struct.unpack("<f", float32_min)[0]),
929 np.float32(struct.unpack("<f", float32_max)[0]),
930 ),
931 "d": (
932 np.float64(struct.unpack("<d", float64_min)[0]),
933 np.float64(struct.unpack("<d", float64_max)[0]),
934 ),
935 }
937 self.OLD_TYPE_MAPPING = {
938 98: 251, # byte
939 105: 252, # int
940 108: 253, # long
941 102: 254, # float
942 100: 255, # double
943 }
945 # These missing values are the generic '.' in Stata, and are used
946 # to replace nans
947 self.MISSING_VALUES: dict[str, int | np.float32 | np.float64] = {
948 "b": 101,
949 "h": 32741,
950 "l": 2147483621,
951 "f": np.float32(struct.unpack("<f", b"\x00\x00\x00\x7f")[0]),
952 "d": np.float64(
953 struct.unpack("<d", b"\x00\x00\x00\x00\x00\x00\xe0\x7f")[0]
954 ),
955 }
956 self.NUMPY_TYPE_MAP = {
957 "b": "i1",
958 "h": "i2",
959 "l": "i4",
960 "f": "f4",
961 "d": "f8",
962 "Q": "u8",
963 }
965 # Reserved words cannot be used as variable names
966 self.RESERVED_WORDS = {
967 "aggregate",
968 "array",
969 "boolean",
970 "break",
971 "byte",
972 "case",
973 "catch",
974 "class",
975 "colvector",
976 "complex",
977 "const",
978 "continue",
979 "default",
980 "delegate",
981 "delete",
982 "do",
983 "double",
984 "else",
985 "eltypedef",
986 "end",
987 "enum",
988 "explicit",
989 "export",
990 "external",
991 "float",
992 "for",
993 "friend",
994 "function",
995 "global",
996 "goto",
997 "if",
998 "inline",
999 "int",
1000 "local",
1001 "long",
1002 "NULL",
1003 "pragma",
1004 "protected",
1005 "quad",
1006 "rowvector",
1007 "short",
1008 "typedef",
1009 "typename",
1010 "virtual",
1011 "_all",
1012 "_N",
1013 "_skip",
1014 "_b",
1015 "_pi",
1016 "str#",
1017 "in",
1018 "_pred",
1019 "strL",
1020 "_coef",
1021 "_rc",
1022 "using",
1023 "_cons",
1024 "_se",
1025 "with",
1026 "_n",
1027 }
1030@set_module("pandas.api.typing")
1031class StataReader(StataParser, abc.Iterator[DataFrame]):
1032 """
1033 Class for reading Stata dta files.
1035 Parameters
1036 ----------
1037 path_or_buf : path (string), buffer or path object
1038 string, pathlib.Path or object
1039 implementing a binary read() functions.
1040 convert_dates : bool, default True
1041 Convert date variables to DataFrame time values.
1042 convert_categoricals : bool, default True
1043 Read value labels and convert columns to Categorical/Factor variables.
1044 index_col : str, optional
1045 Column to set as index.
1046 convert_missing : bool, default False
1047 Flag indicating whether to convert missing values to their Stata
1048 representations. If False, missing values are replaced with nan.
1049 If True, columns containing missing values are returned with
1050 object data types and missing values are represented by
1051 StataMissingValue objects.
1052 preserve_dtypes : bool, default True
1053 Preserve Stata datatypes. If False, numeric data are upcast to pandas
1054 default types for foreign data (float64 or int64).
1055 columns : list or None
1056 Columns to retain. Columns will be returned in the given order. None
1057 returns all columns.
1058 order_categoricals : bool, default True
1059 Flag indicating whether converted categorical data are ordered.
1060 chunksize : int, default None
1061 Return StataReader object for iterations, returns chunks with
1062 given number of lines.
1063 compression : str or dict, default 'infer'
1064 For on-the-fly decompression of on-disk data. If 'infer' and '%s' is
1065 path-like, then detect compression from the following extensions: '.gz',
1066 '.bz2', '.zip', '.xz', '.zst', '.tar', '.tar.gz', '.tar.xz' or '.tar.bz2'
1067 (otherwise no compression).
1068 If using 'zip' or 'tar', the ZIP file must contain only one data file
1069 to be read in.
1070 Set to ``None`` for no decompression.
1071 Can also be a dict with key ``'method'`` set
1072 to one of ``'zip'``, ``'gzip'``, ``'bz2'``, ``'zstd'``, ``'xz'``, ``'tar'`` and
1073 other key-value pairs are forwarded to
1074 ``zipfile.ZipFile``, ``gzip.GzipFile``,
1075 ``bz2.BZ2File``, ``zstandard.ZstdDecompressor``, ``lzma.LZMAFile`` or
1076 ``tarfile.TarFile``, respectively.
1077 As an example, the following could be passed for Zstandard decompression using a
1078 custom compression dictionary:
1079 ``compression={'method': 'zstd', 'dict_data': my_compression_dict}``.
1080 storage_options : dict, optional
1081 Extra options that make sense for a particular storage connection, e.g.
1082 host, port, username, password, etc. For HTTP(S) URLs the key-value pairs
1083 are forwarded to ``urllib.request.Request`` as header options. For other
1084 URLs (e.g. starting with "s3://", and "gcs://") the key-value pairs are
1085 forwarded to ``fsspec.open``. Please see ``fsspec`` and ``urllib`` for more
1086 details, and for more examples on storage options refer `here
1087 <https://pandas.pydata.org/docs/user_guide/io.html?
1088 highlight=storage_options#reading-writing-remote-files>`_.
1090 Notes
1091 -----
1092 Categorical variables read through an iterator may not have the same
1093 categories and dtype. This occurs when a variable stored in a DTA
1094 file is associated to an incomplete set of value labels that only
1095 label a strict subset of the values.
1096 """
1098 _path_or_buf: IO[bytes] # Will be assigned by `_open_file`.
1100 def __init__(
1101 self,
1102 path_or_buf: FilePath | ReadBuffer[bytes],
1103 convert_dates: bool = True,
1104 convert_categoricals: bool = True,
1105 index_col: str | None = None,
1106 convert_missing: bool = False,
1107 preserve_dtypes: bool = True,
1108 columns: Sequence[str] | None = None,
1109 order_categoricals: bool = True,
1110 chunksize: int | None = None,
1111 compression: CompressionOptions = "infer",
1112 storage_options: StorageOptions | None = None,
1113 ) -> None:
1114 super().__init__()
1116 # Arguments to the reader (can be temporarily overridden in
1117 # calls to read).
1118 self._convert_dates = convert_dates
1119 self._convert_categoricals = convert_categoricals
1120 self._index_col = index_col
1121 self._convert_missing = convert_missing
1122 self._preserve_dtypes = preserve_dtypes
1123 self._columns = columns
1124 self._order_categoricals = order_categoricals
1125 self._original_path_or_buf = path_or_buf
1126 self._compression = compression
1127 self._storage_options = storage_options
1128 self._encoding = ""
1129 self._chunksize = chunksize
1130 self._using_iterator = False
1131 self._entered = False
1132 if self._chunksize is None:
1133 self._chunksize = 1
1134 elif not isinstance(chunksize, int) or chunksize <= 0:
1135 raise ValueError("chunksize must be a positive integer when set.")
1137 # State variables for the file
1138 self._close_file: Callable[[], None] | None = None
1139 self._column_selector_set = False
1140 self._value_label_dict: dict[str, dict[int, str]] = {}
1141 self._value_labels_read = False
1142 self._dtype: np.dtype | None = None
1143 self._lines_read = 0
1145 self._native_byteorder = _set_endianness(sys.byteorder)
1147 def _ensure_open(self) -> None:
1148 """
1149 Ensure the file has been opened and its header data read.
1150 """
1151 if not hasattr(self, "_path_or_buf"):
1152 self._open_file()
1154 def _open_file(self) -> None:
1155 """
1156 Open the file (with compression options, etc.), and read header information.
1157 """
1158 if not self._entered:
1159 warnings.warn(
1160 "StataReader is being used without using a context manager. "
1161 "Using StataReader as a context manager is the only supported method.",
1162 ResourceWarning,
1163 stacklevel=find_stack_level(),
1164 )
1165 handles = get_handle(
1166 self._original_path_or_buf,
1167 "rb",
1168 storage_options=self._storage_options,
1169 is_text=False,
1170 compression=self._compression,
1171 )
1172 if hasattr(handles.handle, "seekable") and handles.handle.seekable():
1173 # If the handle is directly seekable, use it without an extra copy.
1174 self._path_or_buf = handles.handle
1175 self._close_file = handles.close
1176 else:
1177 # Copy to memory, and ensure no encoding.
1178 with handles:
1179 self._path_or_buf = BytesIO(handles.handle.read())
1180 self._close_file = self._path_or_buf.close
1182 self._read_header()
1183 self._setup_dtype()
1185 def __enter__(self) -> Self:
1186 """enter context manager"""
1187 self._entered = True
1188 return self
1190 def __exit__(
1191 self,
1192 exc_type: type[BaseException] | None,
1193 exc_value: BaseException | None,
1194 traceback: TracebackType | None,
1195 ) -> None:
1196 if self._close_file:
1197 self._close_file()
1199 def _set_encoding(self) -> None:
1200 """
1201 Set string encoding which depends on file version
1202 """
1203 if self._format_version < 118:
1204 self._encoding = "latin-1"
1205 else:
1206 self._encoding = "utf-8"
1208 def _read_int8(self) -> int:
1209 return struct.unpack("b", self._path_or_buf.read(1))[0]
1211 def _read_uint8(self) -> int:
1212 return struct.unpack("B", self._path_or_buf.read(1))[0]
1214 def _read_uint16(self) -> int:
1215 return struct.unpack(f"{self._byteorder}H", self._path_or_buf.read(2))[0]
1217 def _read_uint32(self) -> int:
1218 return struct.unpack(f"{self._byteorder}I", self._path_or_buf.read(4))[0]
1220 def _read_uint64(self) -> int:
1221 return struct.unpack(f"{self._byteorder}Q", self._path_or_buf.read(8))[0]
1223 def _read_int16(self) -> int:
1224 return struct.unpack(f"{self._byteorder}h", self._path_or_buf.read(2))[0]
1226 def _read_int32(self) -> int:
1227 return struct.unpack(f"{self._byteorder}i", self._path_or_buf.read(4))[0]
1229 def _read_int64(self) -> int:
1230 return struct.unpack(f"{self._byteorder}q", self._path_or_buf.read(8))[0]
1232 def _read_char8(self) -> bytes:
1233 return struct.unpack("c", self._path_or_buf.read(1))[0]
1235 def _read_int16_count(self, count: int) -> tuple[int, ...]:
1236 return struct.unpack(
1237 f"{self._byteorder}{'h' * count}",
1238 self._path_or_buf.read(2 * count),
1239 )
1241 def _read_header(self) -> None:
1242 first_char = self._read_char8()
1243 if first_char == b"<":
1244 self._read_new_header()
1245 else:
1246 self._read_old_header(first_char)
1248 def _read_new_header(self) -> None:
1249 # The first part of the header is common to 117 - 119.
1250 self._path_or_buf.read(27) # stata_dta><header><release>
1251 self._format_version = int(self._path_or_buf.read(3))
1252 if self._format_version not in _new_format_versions:
1253 raise ValueError(_version_error.format(version=self._format_version))
1254 self._set_encoding()
1255 self._path_or_buf.read(21) # </release><byteorder>
1256 self._byteorder = ">" if self._path_or_buf.read(3) == b"MSF" else "<"
1257 self._path_or_buf.read(15) # </byteorder><K>
1258 self._nvar = (
1259 self._read_uint16() if self._format_version <= 118 else self._read_uint32()
1260 )
1261 self._path_or_buf.read(7) # </K><N>
1263 self._nobs = self._get_nobs()
1264 self._path_or_buf.read(11) # </N><label>
1265 self._data_label = self._get_data_label()
1266 self._path_or_buf.read(19) # </label><timestamp>
1267 self._time_stamp = self._get_time_stamp()
1268 self._path_or_buf.read(26) # </timestamp></header><map>
1269 self._path_or_buf.read(8) # 0x0000000000000000
1270 self._path_or_buf.read(8) # position of <map>
1272 self._seek_vartypes = self._read_int64() + 16
1273 self._seek_varnames = self._read_int64() + 10
1274 self._seek_sortlist = self._read_int64() + 10
1275 self._seek_formats = self._read_int64() + 9
1276 self._seek_value_label_names = self._read_int64() + 19
1278 # Requires version-specific treatment
1279 self._seek_variable_labels = self._get_seek_variable_labels()
1281 self._path_or_buf.read(8) # <characteristics>
1282 self._data_location = self._read_int64() + 6
1283 self._seek_strls = self._read_int64() + 7
1284 self._seek_value_labels = self._read_int64() + 14
1286 self._typlist, self._dtyplist = self._get_dtypes(self._seek_vartypes)
1288 self._path_or_buf.seek(self._seek_varnames)
1289 self._varlist = self._get_varlist()
1291 self._path_or_buf.seek(self._seek_sortlist)
1292 self._srtlist = self._read_int16_count(self._nvar + 1)[:-1]
1294 self._path_or_buf.seek(self._seek_formats)
1295 self._fmtlist = self._get_fmtlist()
1297 self._path_or_buf.seek(self._seek_value_label_names)
1298 self._lbllist = self._get_lbllist()
1300 self._path_or_buf.seek(self._seek_variable_labels)
1301 self._variable_labels = self._get_variable_labels()
1303 # Get data type information, works for versions 117-119.
1304 def _get_dtypes(
1305 self, seek_vartypes: int
1306 ) -> tuple[list[int | str], list[str | np.dtype]]:
1307 self._path_or_buf.seek(seek_vartypes)
1308 typlist = []
1309 dtyplist = []
1310 for _ in range(self._nvar):
1311 typ = self._read_uint16()
1312 if typ <= 2045:
1313 typlist.append(typ)
1314 dtyplist.append(str(typ))
1315 else:
1316 try:
1317 typlist.append(self.TYPE_MAP_XML[typ]) # type: ignore[arg-type]
1318 dtyplist.append(self.DTYPE_MAP_XML[typ]) # type: ignore[arg-type]
1319 except KeyError as err:
1320 raise ValueError(f"cannot convert stata types [{typ}]") from err
1322 return typlist, dtyplist # type: ignore[return-value]
1324 def _get_varlist(self) -> list[str]:
1325 # 33 in order formats, 129 in formats 118 and 119
1326 b = 33 if self._format_version < 118 else 129
1327 return [self._decode(self._path_or_buf.read(b)) for _ in range(self._nvar)]
1329 # Returns the format list
1330 def _get_fmtlist(self) -> list[str]:
1331 if self._format_version >= 118:
1332 b = 57
1333 elif self._format_version > 113:
1334 b = 49
1335 elif self._format_version > 104:
1336 b = 12
1337 else:
1338 b = 7
1340 return [self._decode(self._path_or_buf.read(b)) for _ in range(self._nvar)]
1342 # Returns the label list
1343 def _get_lbllist(self) -> list[str]:
1344 if self._format_version >= 118:
1345 b = 129
1346 elif self._format_version > 108:
1347 b = 33
1348 else:
1349 b = 9
1350 return [self._decode(self._path_or_buf.read(b)) for _ in range(self._nvar)]
1352 def _get_variable_labels(self) -> list[str]:
1353 if self._format_version >= 118:
1354 vlblist = [
1355 self._decode(self._path_or_buf.read(321)) for _ in range(self._nvar)
1356 ]
1357 elif self._format_version > 105:
1358 vlblist = [
1359 self._decode(self._path_or_buf.read(81)) for _ in range(self._nvar)
1360 ]
1361 else:
1362 vlblist = [
1363 self._decode(self._path_or_buf.read(32)) for _ in range(self._nvar)
1364 ]
1365 return vlblist
1367 def _get_nobs(self) -> int:
1368 if self._format_version >= 118:
1369 return self._read_uint64()
1370 elif self._format_version >= 103:
1371 return self._read_uint32()
1372 else:
1373 return self._read_uint16()
1375 def _get_data_label(self) -> str:
1376 if self._format_version >= 118:
1377 strlen = self._read_uint16()
1378 return self._decode(self._path_or_buf.read(strlen))
1379 elif self._format_version == 117:
1380 strlen = self._read_int8()
1381 return self._decode(self._path_or_buf.read(strlen))
1382 elif self._format_version > 105:
1383 return self._decode(self._path_or_buf.read(81))
1384 else:
1385 return self._decode(self._path_or_buf.read(32))
1387 def _get_time_stamp(self) -> str:
1388 if self._format_version >= 118:
1389 strlen = self._read_int8()
1390 return self._path_or_buf.read(strlen).decode("utf-8")
1391 elif self._format_version == 117:
1392 strlen = self._read_int8()
1393 return self._decode(self._path_or_buf.read(strlen))
1394 elif self._format_version > 104:
1395 return self._decode(self._path_or_buf.read(18))
1396 else:
1397 raise ValueError
1399 def _get_seek_variable_labels(self) -> int:
1400 if self._format_version == 117:
1401 self._path_or_buf.read(8) # <variable_labels>, throw away
1402 # Stata 117 data files do not follow the described format. This is
1403 # a work around that uses the previous label, 33 bytes for each
1404 # variable, 20 for the closing tag and 17 for the opening tag
1405 return self._seek_value_label_names + (33 * self._nvar) + 20 + 17
1406 elif self._format_version >= 118:
1407 return self._read_int64() + 17
1408 else:
1409 raise ValueError
1411 def _read_old_header(self, first_char: bytes) -> None:
1412 self._format_version = int(first_char[0])
1413 if self._format_version not in _old_format_versions:
1414 raise ValueError(_version_error.format(version=self._format_version))
1415 self._set_encoding()
1416 # Note 102 format will have a zero in this header position, so support
1417 # relies on little-endian being set whenever this value isn't one,
1418 # even though for later releases strictly speaking the value should
1419 # be either one or two to be valid
1420 self._byteorder = ">" if self._read_int8() == 0x1 else "<"
1421 self._filetype = self._read_int8()
1422 self._path_or_buf.read(1) # unused
1424 self._nvar = self._read_uint16()
1425 self._nobs = self._get_nobs()
1427 self._data_label = self._get_data_label()
1429 if self._format_version >= 105:
1430 self._time_stamp = self._get_time_stamp()
1432 # descriptors
1433 if self._format_version >= 111:
1434 typlist = [int(c) for c in self._path_or_buf.read(self._nvar)]
1435 else:
1436 buf = self._path_or_buf.read(self._nvar)
1437 typlistb = np.frombuffer(buf, dtype=np.uint8)
1438 typlist = []
1439 for tp in typlistb:
1440 if tp in self.OLD_TYPE_MAPPING:
1441 typlist.append(self.OLD_TYPE_MAPPING[tp]) # type: ignore[index]
1442 else:
1443 typlist.append(tp - 127) # type: ignore[arg-type]
1445 try:
1446 self._typlist = [self.TYPE_MAP[typ] for typ in typlist]
1447 except ValueError as err:
1448 invalid_types = ",".join([str(x) for x in typlist])
1449 raise ValueError(f"cannot convert stata types [{invalid_types}]") from err
1450 try:
1451 self._dtyplist = [self.DTYPE_MAP[typ] for typ in typlist]
1452 except ValueError as err:
1453 invalid_dtypes = ",".join([str(x) for x in typlist])
1454 raise ValueError(f"cannot convert stata dtypes [{invalid_dtypes}]") from err
1456 if self._format_version > 108:
1457 self._varlist = [
1458 self._decode(self._path_or_buf.read(33)) for _ in range(self._nvar)
1459 ]
1460 else:
1461 self._varlist = [
1462 self._decode(self._path_or_buf.read(9)) for _ in range(self._nvar)
1463 ]
1464 self._srtlist = self._read_int16_count(self._nvar + 1)[:-1]
1466 self._fmtlist = self._get_fmtlist()
1468 self._lbllist = self._get_lbllist()
1470 self._variable_labels = self._get_variable_labels()
1472 # ignore expansion fields (Format 105 and later)
1473 # When reading, read five bytes; the last four bytes now tell you
1474 # the size of the next read, which you discard. You then continue
1475 # like this until you read 5 bytes of zeros.
1477 if self._format_version > 104:
1478 while True:
1479 data_type = self._read_int8()
1480 if self._format_version > 108:
1481 data_len = self._read_int32()
1482 else:
1483 data_len = self._read_int16()
1484 if data_type == 0:
1485 break
1486 self._path_or_buf.read(data_len)
1488 # necessary data to continue parsing
1489 self._data_location = self._path_or_buf.tell()
1491 def _setup_dtype(self) -> np.dtype:
1492 """Map between numpy and state dtypes"""
1493 if self._dtype is not None:
1494 return self._dtype
1496 dtypes = [] # Convert struct data types to numpy data type
1497 for i, typ in enumerate(self._typlist):
1498 if typ in self.NUMPY_TYPE_MAP:
1499 dtypes.append((f"s{i}", f"{self._byteorder}{self.NUMPY_TYPE_MAP[typ]}"))
1500 else:
1501 dtypes.append((f"s{i}", f"S{typ}"))
1502 self._dtype = np.dtype(dtypes)
1504 return self._dtype
1506 def _decode(self, s: bytes) -> str:
1507 # have bytes not strings, so must decode
1508 s = s.partition(b"\0")[0]
1509 try:
1510 return s.decode(self._encoding)
1511 except UnicodeDecodeError:
1512 # GH 25960, fallback to handle incorrect format produced when 117
1513 # files are converted to 118 files in Stata
1514 encoding = self._encoding
1515 msg = f"""
1516One or more strings in the dta file could not be decoded using {encoding}, and
1517so the fallback encoding of latin-1 is being used. This can happen when a file
1518has been incorrectly encoded by Stata or some other software. You should verify
1519the string values returned are correct."""
1520 warnings.warn(
1521 msg,
1522 UnicodeWarning,
1523 stacklevel=find_stack_level(),
1524 )
1525 return s.decode("latin-1")
1527 def _read_new_value_labels(self) -> None:
1528 """Reads value labels with variable length strings (108 and later format)"""
1529 if self._format_version >= 117:
1530 self._path_or_buf.seek(self._seek_value_labels)
1531 else:
1532 assert self._dtype is not None
1533 offset = self._nobs * self._dtype.itemsize
1534 self._path_or_buf.seek(self._data_location + offset)
1536 while True:
1537 if self._format_version >= 117:
1538 if self._path_or_buf.read(5) == b"</val": # <lbl>
1539 break # end of value label table
1541 slength = self._path_or_buf.read(4)
1542 if not slength:
1543 break # end of value label table (format < 117), or end-of-file
1544 if self._format_version == 108:
1545 labname = self._decode(self._path_or_buf.read(9))
1546 elif self._format_version <= 117:
1547 labname = self._decode(self._path_or_buf.read(33))
1548 else:
1549 labname = self._decode(self._path_or_buf.read(129))
1550 self._path_or_buf.read(3) # padding
1552 n = self._read_uint32()
1553 txtlen = self._read_uint32()
1554 off = np.frombuffer(
1555 self._path_or_buf.read(4 * n), dtype=f"{self._byteorder}i4", count=n
1556 )
1557 val = np.frombuffer(
1558 self._path_or_buf.read(4 * n), dtype=f"{self._byteorder}i4", count=n
1559 )
1560 ii = np.argsort(off)
1561 off = off[ii]
1562 val = val[ii]
1563 txt = self._path_or_buf.read(txtlen)
1564 self._value_label_dict[labname] = {}
1565 for i in range(n):
1566 end = off[i + 1] if i < n - 1 else txtlen
1567 self._value_label_dict[labname][val[i]] = self._decode(
1568 txt[off[i] : end]
1569 )
1571 if self._format_version >= 117:
1572 self._path_or_buf.read(6) # </lbl>
1574 def _read_old_value_labels(self) -> None:
1575 """Reads value labels with fixed-length strings (105 and earlier format)"""
1576 assert self._dtype is not None
1577 offset = self._nobs * self._dtype.itemsize
1578 self._path_or_buf.seek(self._data_location + offset)
1580 while True:
1581 if not self._path_or_buf.read(2):
1582 # end-of-file may have been reached, if so stop here
1583 break
1585 # otherwise back up and read again, taking byteorder into account
1586 self._path_or_buf.seek(-2, os.SEEK_CUR)
1587 n = self._read_uint16()
1588 labname = self._decode(self._path_or_buf.read(9))
1589 self._path_or_buf.read(1) # padding
1590 codes = np.frombuffer(
1591 self._path_or_buf.read(2 * n), dtype=f"{self._byteorder}i2", count=n
1592 )
1593 self._value_label_dict[labname] = {}
1594 for i in range(n):
1595 self._value_label_dict[labname][codes[i]] = self._decode(
1596 self._path_or_buf.read(8)
1597 )
1599 def _read_value_labels(self) -> None:
1600 self._ensure_open()
1601 if self._value_labels_read:
1602 # Don't read twice
1603 return
1605 if self._format_version >= 108:
1606 self._read_new_value_labels()
1607 else:
1608 self._read_old_value_labels()
1609 self._value_labels_read = True
1611 def _read_strls(self) -> None:
1612 self._path_or_buf.seek(self._seek_strls)
1613 # Wrap v_o in a string to allow uint64 values as keys on 32bit OS
1614 self.GSO = {"0": ""}
1615 while True:
1616 if self._path_or_buf.read(3) != b"GSO":
1617 break
1619 if self._format_version == 117:
1620 v_o = self._read_uint64()
1621 else:
1622 buf = self._path_or_buf.read(12)
1623 # Only tested on little endian machine.
1624 v_size = 2 if self._format_version == 118 else 3
1625 if self._byteorder == "<":
1626 buf = buf[0:v_size] + buf[4 : (12 - v_size)]
1627 else:
1628 buf = buf[4 - v_size : 4] + buf[(4 + v_size) :]
1629 v_o = struct.unpack(f"{self._byteorder}Q", buf)[0]
1630 typ = self._read_uint8()
1631 length = self._read_uint32()
1632 va = self._path_or_buf.read(length)
1633 if typ == 130:
1634 decoded_va = va[0:-1].decode(self._encoding)
1635 else:
1636 # Stata says typ 129 can be binary, so use str
1637 decoded_va = str(va)
1638 # Wrap v_o in a string to allow uint64 values as keys on 32bit OS
1639 self.GSO[str(v_o)] = decoded_va
1641 def __next__(self) -> DataFrame:
1642 self._using_iterator = True
1643 return self.read(nrows=self._chunksize)
1645 def get_chunk(self, size: int | None = None) -> DataFrame:
1646 """
1647 Reads lines from Stata file and returns as dataframe
1649 Parameters
1650 ----------
1651 size : int, defaults to None
1652 Number of lines to read. If None, reads whole file.
1654 Returns
1655 -------
1656 DataFrame
1657 """
1658 if size is None:
1659 size = self._chunksize
1660 return self.read(nrows=size)
1662 def read(
1663 self,
1664 nrows: int | None = None,
1665 convert_dates: bool | None = None,
1666 convert_categoricals: bool | None = None,
1667 index_col: str | None = None,
1668 convert_missing: bool | None = None,
1669 preserve_dtypes: bool | None = None,
1670 columns: Sequence[str] | None = None,
1671 order_categoricals: bool | None = None,
1672 ) -> DataFrame:
1673 """
1674 Reads observations from Stata file, converting them into a dataframe
1676 Parameters
1677 ----------
1678 nrows : int
1679 Number of lines to read from data file, if None read whole file.
1680 convert_dates : bool, default True
1681 Convert date variables to DataFrame time values.
1682 convert_categoricals : bool, default True
1683 Read value labels and convert columns to Categorical/Factor variables.
1684 index_col : str, optional
1685 Column to set as index.
1686 convert_missing : bool, default False
1687 Flag indicating whether to convert missing values to their Stata
1688 representations. If False, missing values are replaced with nan.
1689 If True, columns containing missing values are returned with
1690 object data types and missing values are represented by
1691 StataMissingValue objects.
1692 preserve_dtypes : bool, default True
1693 Preserve Stata datatypes. If False, numeric data are upcast to pandas
1694 default types for foreign data (float64 or int64).
1695 columns : list or None
1696 Columns to retain. Columns will be returned in the given order. None
1697 returns all columns.
1698 order_categoricals : bool, default True
1699 Flag indicating whether converted categorical data are ordered.
1701 Returns
1702 -------
1703 DataFrame
1704 """
1705 self._ensure_open()
1707 # Handle options
1708 if convert_dates is None:
1709 convert_dates = self._convert_dates
1710 if convert_categoricals is None:
1711 convert_categoricals = self._convert_categoricals
1712 if convert_missing is None:
1713 convert_missing = self._convert_missing
1714 if preserve_dtypes is None:
1715 preserve_dtypes = self._preserve_dtypes
1716 if columns is None:
1717 columns = self._columns
1718 if order_categoricals is None:
1719 order_categoricals = self._order_categoricals
1720 if index_col is None:
1721 index_col = self._index_col
1722 if nrows is None:
1723 nrows = self._nobs
1725 # Handle empty file or chunk. If reading incrementally raise
1726 # StopIteration. If reading the whole thing return an empty
1727 # data frame.
1728 if (self._nobs == 0) and nrows == 0:
1729 data = DataFrame(columns=self._varlist)
1730 # Apply dtypes correctly
1731 for i, col in enumerate(data.columns):
1732 dt = self._dtyplist[i]
1733 if isinstance(dt, np.dtype):
1734 if dt.char != "S":
1735 data[col] = data[col].astype(dt)
1736 if columns is not None:
1737 data = self._do_select_columns(data, columns)
1738 return data
1740 if (self._format_version >= 117) and (not self._value_labels_read):
1741 self._read_strls()
1743 # Read data
1744 assert self._dtype is not None
1745 dtype = self._dtype
1746 max_read_len = (self._nobs - self._lines_read) * dtype.itemsize
1747 read_len = nrows * dtype.itemsize
1748 read_len = min(read_len, max_read_len)
1749 if read_len <= 0:
1750 # Iterator has finished, should never be here unless
1751 # we are reading the file incrementally
1752 if convert_categoricals:
1753 self._read_value_labels()
1754 raise StopIteration
1755 offset = self._lines_read * dtype.itemsize
1756 self._path_or_buf.seek(self._data_location + offset)
1757 read_lines = min(nrows, self._nobs - self._lines_read)
1758 raw_data = np.frombuffer(
1759 self._path_or_buf.read(read_len), dtype=dtype, count=read_lines
1760 )
1762 self._lines_read += read_lines
1764 # if necessary, swap the byte order to native here
1765 if self._byteorder != self._native_byteorder:
1766 raw_data = raw_data.byteswap().view(raw_data.dtype.newbyteorder())
1768 if convert_categoricals:
1769 self._read_value_labels()
1771 if len(raw_data) == 0:
1772 data = DataFrame(columns=self._varlist)
1773 else:
1774 data = DataFrame.from_records(raw_data)
1775 data.columns = Index(self._varlist)
1777 # If index is not specified, use actual row number rather than
1778 # restarting at 0 for each chunk.
1779 if index_col is None:
1780 data.index = RangeIndex(
1781 self._lines_read - read_lines, self._lines_read
1782 ) # set attr instead of set_index to avoid copy
1784 if columns is not None:
1785 data = self._do_select_columns(data, columns)
1787 # Decode strings
1788 for col, typ in zip(data, self._typlist, strict=True):
1789 if isinstance(typ, int):
1790 data[col] = data[col].apply(self._decode)
1792 data = self._insert_strls(data)
1794 # Convert columns (if needed) to match input type
1795 valid_dtypes = [i for i, dtyp in enumerate(self._dtyplist) if dtyp is not None]
1796 object_type = np.dtype(object)
1797 for idx in valid_dtypes:
1798 dtype = data.iloc[:, idx].dtype
1799 if dtype not in (object_type, self._dtyplist[idx]):
1800 data.isetitem(idx, data.iloc[:, idx].astype(dtype))
1802 data = self._do_convert_missing(data, convert_missing)
1804 if convert_dates:
1805 for i, fmt in enumerate(self._fmtlist):
1806 if any(fmt.startswith(date_fmt) for date_fmt in _date_formats):
1807 data.isetitem(
1808 i, _stata_elapsed_date_to_datetime_vec(data.iloc[:, i], fmt)
1809 )
1811 if convert_categoricals:
1812 data = self._do_convert_categoricals(
1813 data, self._value_label_dict, self._lbllist, order_categoricals
1814 )
1816 if not preserve_dtypes:
1817 retyped_data = []
1818 convert = False
1819 for col in data:
1820 dtype = data[col].dtype
1821 if dtype in (np.dtype(np.float16), np.dtype(np.float32)):
1822 dtype = np.dtype(np.float64)
1823 convert = True
1824 elif dtype in (
1825 np.dtype(np.int8),
1826 np.dtype(np.int16),
1827 np.dtype(np.int32),
1828 ):
1829 dtype = np.dtype(np.int64)
1830 convert = True
1831 retyped_data.append((col, data[col].astype(dtype)))
1832 if convert:
1833 data = DataFrame.from_dict(dict(retyped_data))
1835 if index_col is not None:
1836 data = data.set_index(data.pop(index_col))
1838 return data
1840 def _do_convert_missing(self, data: DataFrame, convert_missing: bool) -> DataFrame:
1841 # missing code for double was different in version 105 and prior
1842 old_missingdouble = float.fromhex("0x1.0p333")
1844 # Check for missing values, and replace if found
1845 replacements = {}
1846 for i in range(len(data.columns)):
1847 fmt = self._typlist[i]
1848 # recode instances of the old missing code to the currently used value
1849 if self._format_version <= 105 and fmt == "d":
1850 data.iloc[:, i] = data.iloc[:, i].replace(
1851 old_missingdouble, self.MISSING_VALUES["d"]
1852 )
1854 if self._format_version <= 111:
1855 if fmt not in self.OLD_VALID_RANGE:
1856 continue
1858 nmin, nmax = self.OLD_VALID_RANGE[fmt]
1859 else:
1860 if fmt not in self.VALID_RANGE:
1861 continue
1863 nmin, nmax = self.VALID_RANGE[fmt]
1864 series = data.iloc[:, i]
1866 # appreciably faster to do this with ndarray instead of Series
1867 svals = series._values
1868 missing = (svals < nmin) | (svals > nmax)
1870 if not missing.any():
1871 continue
1873 if convert_missing: # Replacement follows Stata notation
1874 missing_loc = np.nonzero(np.asarray(missing))[0]
1875 umissing, umissing_loc = np.unique(series[missing], return_inverse=True)
1876 replacement = Series(series, dtype=object)
1877 for j, um in enumerate(umissing):
1878 if self._format_version <= 111:
1879 missing_value = StataMissingValue(
1880 float(self.MISSING_VALUES[fmt])
1881 )
1882 else:
1883 missing_value = StataMissingValue(um)
1885 loc = missing_loc[umissing_loc == j]
1886 replacement.iloc[loc] = missing_value
1887 else: # All replacements are identical
1888 dtype = series.dtype
1889 if dtype not in (np.float32, np.float64):
1890 dtype = np.float64
1891 replacement = Series(series, dtype=dtype)
1892 # Note: operating on ._values is much faster than directly
1893 # TODO: can we fix that?
1894 replacement._values[missing] = np.nan
1895 replacements[i] = replacement
1896 if replacements:
1897 for idx, value in replacements.items():
1898 data.isetitem(idx, value)
1899 return data
1901 def _insert_strls(self, data: DataFrame) -> DataFrame:
1902 if not hasattr(self, "GSO") or len(self.GSO) == 0:
1903 return data
1904 for i, typ in enumerate(self._typlist):
1905 if typ != "Q":
1906 continue
1907 # Wrap v_o in a string to allow uint64 values as keys on 32bit OS
1908 data.isetitem(i, [self.GSO[str(k)] for k in data.iloc[:, i]])
1909 return data
1911 def _do_select_columns(self, data: DataFrame, columns: Sequence[str]) -> DataFrame:
1912 if not self._column_selector_set:
1913 column_set = set(columns)
1914 if len(column_set) != len(columns):
1915 raise ValueError("columns contains duplicate entries")
1916 unmatched = column_set.difference(data.columns)
1917 if unmatched:
1918 joined = ", ".join(list(unmatched))
1919 raise ValueError(
1920 "The following columns were not "
1921 f"found in the Stata data set: {joined}"
1922 )
1923 # Copy information for retained columns for later processing
1924 dtyplist = []
1925 typlist = []
1926 fmtlist = []
1927 lbllist = []
1928 for col in columns:
1929 i = data.columns.get_loc(col)
1930 dtyplist.append(self._dtyplist[i])
1931 typlist.append(self._typlist[i])
1932 fmtlist.append(self._fmtlist[i])
1933 lbllist.append(self._lbllist[i])
1935 self._dtyplist = dtyplist # type: ignore[assignment]
1936 self._typlist = typlist # type: ignore[assignment]
1937 self._fmtlist = fmtlist # type: ignore[assignment]
1938 self._lbllist = lbllist # type: ignore[assignment]
1939 self._column_selector_set = True
1941 return data[columns]
1943 def _do_convert_categoricals(
1944 self,
1945 data: DataFrame,
1946 value_label_dict: dict[str, dict[int, str]],
1947 lbllist: Sequence[str],
1948 order_categoricals: bool,
1949 ) -> DataFrame:
1950 """
1951 Converts categorical columns to Categorical type.
1952 """
1953 if not value_label_dict:
1954 return data
1955 cat_converted_data = []
1956 for col, label in zip(data, lbllist, strict=True):
1957 if label in value_label_dict:
1958 # Explicit call with ordered=True
1959 vl = value_label_dict[label]
1960 keys = np.array(list(vl.keys()))
1961 column = data[col]
1962 key_matches = column.isin(keys)
1963 if self._using_iterator and key_matches.all():
1964 initial_categories: np.ndarray | None = keys
1965 # If all categories are in the keys and we are iterating,
1966 # use the same keys for all chunks. If some are missing
1967 # value labels, then we will fall back to the categories
1968 # varying across chunks.
1969 else:
1970 if self._using_iterator:
1971 # warn is using an iterator
1972 warnings.warn(
1973 categorical_conversion_warning,
1974 CategoricalConversionWarning,
1975 stacklevel=find_stack_level(),
1976 )
1977 initial_categories = None
1978 cat_data = Categorical(
1979 column, categories=initial_categories, ordered=order_categoricals
1980 )
1981 if initial_categories is None:
1982 # If None here, then we need to match the cats in the Categorical
1983 categories = []
1984 for category in cat_data.categories:
1985 if category in vl:
1986 categories.append(vl[category])
1987 else:
1988 categories.append(category)
1989 else:
1990 # If all cats are matched, we can use the values
1991 categories = list(vl.values())
1992 try:
1993 # Try to catch duplicate categories
1994 # TODO: if we get a non-copying rename_categories, use that
1995 cat_data = cat_data.rename_categories(categories)
1996 except ValueError as err:
1997 vc = Series(categories, copy=False).value_counts()
1998 repeated_cats = list(vc.index[vc > 1])
1999 repeats = "-" * 80 + "\n" + "\n".join(repeated_cats)
2000 # GH 25772
2001 msg = f"""
2002Value labels for column {col} are not unique. These cannot be converted to
2003pandas categoricals.
2005Either read the file with `convert_categoricals` set to False or use the
2006low level interface in `StataReader` to separately read the values and the
2007value_labels. This column's labels are stored under the key '{label}' in
2008`StataReader.value_labels()`.
2010The repeated labels are:
2011{repeats}
2012"""
2013 raise ValueError(msg) from err
2014 # TODO: is the next line needed above in the data(...) method?
2015 cat_series = Series(cat_data, index=data.index, copy=False)
2016 cat_converted_data.append((col, cat_series))
2017 else:
2018 cat_converted_data.append((col, data[col]))
2019 data = DataFrame(dict(cat_converted_data), copy=False)
2020 return data
2022 @property
2023 def data_label(self) -> str:
2024 """
2025 Return data label of Stata file.
2027 The data label is a descriptive string associated with the dataset
2028 stored in the Stata file. This property provides access to that
2029 label, if one is present.
2031 See Also
2032 --------
2033 io.stata.StataReader.variable_labels : Return a dict associating each variable
2034 name with corresponding label.
2035 DataFrame.to_stata : Export DataFrame object to Stata dta format.
2037 Examples
2038 --------
2039 >>> df = pd.DataFrame([(1,)], columns=["variable"])
2040 >>> time_stamp = pd.Timestamp(2000, 2, 29, 14, 21)
2041 >>> data_label = "This is a data file."
2042 >>> path = "/My_path/filename.dta"
2043 >>> df.to_stata(
2044 ... path,
2045 ... time_stamp=time_stamp, # doctest: +SKIP
2046 ... data_label=data_label, # doctest: +SKIP
2047 ... version=None,
2048 ... ) # doctest: +SKIP
2049 >>> with pd.io.stata.StataReader(path) as reader: # doctest: +SKIP
2050 ... print(reader.data_label) # doctest: +SKIP
2051 This is a data file.
2052 """
2053 self._ensure_open()
2054 return self._data_label
2056 @property
2057 def time_stamp(self) -> str:
2058 """
2059 Return time stamp of Stata file.
2060 """
2061 self._ensure_open()
2062 return self._time_stamp
2064 def variable_labels(self) -> dict[str, str]:
2065 """
2066 Return a dict associating each variable name with corresponding label.
2068 This method retrieves variable labels from a Stata file. Variable labels are
2069 mappings between variable names and their corresponding descriptive labels
2070 in a Stata dataset.
2072 Returns
2073 -------
2074 dict
2075 A python dictionary.
2077 See Also
2078 --------
2079 read_stata : Read Stata file into DataFrame.
2080 DataFrame.to_stata : Export DataFrame object to Stata dta format.
2082 Examples
2083 --------
2084 >>> df = pd.DataFrame([[1, 2], [3, 4]], columns=["col_1", "col_2"])
2085 >>> time_stamp = pd.Timestamp(2000, 2, 29, 14, 21)
2086 >>> path = "/My_path/filename.dta"
2087 >>> variable_labels = {"col_1": "This is an example"}
2088 >>> df.to_stata(
2089 ... path,
2090 ... time_stamp=time_stamp, # doctest: +SKIP
2091 ... variable_labels=variable_labels,
2092 ... version=None,
2093 ... ) # doctest: +SKIP
2094 >>> with pd.io.stata.StataReader(path) as reader: # doctest: +SKIP
2095 ... print(reader.variable_labels()) # doctest: +SKIP
2096 {'index': '', 'col_1': 'This is an example', 'col_2': ''}
2097 >>> pd.read_stata(path) # doctest: +SKIP
2098 index col_1 col_2
2099 0 0 1 2
2100 1 1 3 4
2101 """
2102 self._ensure_open()
2103 return dict(zip(self._varlist, self._variable_labels, strict=True))
2105 def value_labels(self) -> dict[str, dict[int, str]]:
2106 """
2107 Return a nested dict mapping each value label name to its values and labels.
2109 This method retrieves the value labels from a Stata file. Value labels are
2110 mappings between the coded values and their corresponding descriptive labels
2111 in a Stata dataset.
2113 Returns
2114 -------
2115 dict
2116 A python dictionary keyed by value label name, which need not match a
2117 variable name. Several variables can share one value label set.
2119 See Also
2120 --------
2121 read_stata : Read Stata file into DataFrame.
2122 DataFrame.to_stata : Export DataFrame object to Stata dta format.
2124 Examples
2125 --------
2126 >>> df = pd.DataFrame([[1, 2], [3, 4]], columns=["col_1", "col_2"])
2127 >>> time_stamp = pd.Timestamp(2000, 2, 29, 14, 21)
2128 >>> path = "/My_path/filename.dta"
2129 >>> value_labels = {"col_1": {3: "x"}}
2130 >>> df.to_stata(
2131 ... path,
2132 ... time_stamp=time_stamp, # doctest: +SKIP
2133 ... value_labels=value_labels,
2134 ... version=None,
2135 ... ) # doctest: +SKIP
2136 >>> with pd.io.stata.StataReader(path) as reader: # doctest: +SKIP
2137 ... print(reader.value_labels()) # doctest: +SKIP
2138 {'col_1': {3: 'x'}}
2139 >>> pd.read_stata(path) # doctest: +SKIP
2140 index col_1 col_2
2141 0 0 1 2
2142 1 1 x 4
2143 """
2144 if not self._value_labels_read:
2145 self._read_value_labels()
2147 return self._value_label_dict
2150@set_module("pandas")
2151def read_stata(
2152 filepath_or_buffer: FilePath | ReadBuffer[bytes],
2153 *,
2154 convert_dates: bool = True,
2155 convert_categoricals: bool = True,
2156 index_col: str | None = None,
2157 convert_missing: bool = False,
2158 preserve_dtypes: bool = True,
2159 columns: Sequence[str] | None = None,
2160 order_categoricals: bool = True,
2161 chunksize: int | None = None,
2162 iterator: bool = False,
2163 compression: CompressionOptions = "infer",
2164 storage_options: StorageOptions | None = None,
2165) -> DataFrame | StataReader:
2166 """
2167 Read Stata file into DataFrame.
2169 This function reads ``.dta`` files produced by Stata, with support for
2170 converting date variables, categorical data, and missing value
2171 representations.
2173 Parameters
2174 ----------
2175 filepath_or_buffer : str, path object or file-like object
2176 Any valid string path is acceptable. The string could be a URL. Valid
2177 URL schemes include http, ftp, s3, and file. For file URLs, a host is
2178 expected. A local file could be: ``file://localhost/path/to/table.dta``.
2180 Certain URL schemes may require additional packages. For example, S3
2181 URLs require the ``s3fs`` library. See
2182 :ref:`install.optional_dependencies` for a full list.
2184 If you want to pass in a path object, pandas accepts any ``os.PathLike``.
2186 By file-like object, we refer to objects with a ``read()`` method,
2187 such as a file handle (e.g. via builtin ``open`` function)
2188 or ``StringIO``.
2189 convert_dates : bool, default True
2190 Convert date variables to DataFrame time values.
2191 convert_categoricals : bool, default True
2192 Read value labels and convert columns to Categorical/Factor variables.
2193 index_col : str, optional
2194 Column to set as index.
2195 convert_missing : bool, default False
2196 Flag indicating whether to convert missing values to their Stata
2197 representations. If False, missing values are replaced with nan.
2198 If True, columns containing missing values are returned with
2199 object data types and missing values are represented by
2200 StataMissingValue objects.
2201 preserve_dtypes : bool, default True
2202 Preserve Stata datatypes. If False, numeric data are upcast to pandas
2203 default types for foreign data (float64 or int64).
2204 columns : list or None
2205 Columns to retain. Columns will be returned in the given order. None
2206 returns all columns.
2207 order_categoricals : bool, default True
2208 Flag indicating whether converted categorical data are ordered.
2209 chunksize : int, default None
2210 Return StataReader object for iterations, returns chunks with
2211 given number of lines.
2212 iterator : bool, default False
2213 Return StataReader object.
2214 compression : str or dict, default 'infer'
2215 For on-the-fly decompression of on-disk data. If 'infer' and
2216 'filepath_or_buffer' is path-like, then detect compression from the
2217 following extensions: '.gz', '.bz2', '.zip', '.xz', '.zst', '.tar',
2218 '.tar.gz', '.tar.xz' or '.tar.bz2' (otherwise no compression).
2219 If using 'zip' or 'tar', the ZIP file must contain only one
2220 data file to be read in. Set to ``None`` for no decompression.
2221 Can also be a dict with key ``'method'`` set to one of
2222 {``'zip'``, ``'gzip'``, ``'bz2'``, ``'zstd'``, ``'xz'``, ``'tar'``} and
2223 other key-value pairs are forwarded to
2224 ``zipfile.ZipFile``, ``gzip.GzipFile``,
2225 ``bz2.BZ2File``, ``zstandard.ZstdDecompressor``, ``lzma.LZMAFile`` or
2226 ``tarfile.TarFile``, respectively.
2227 As an example, the following could be passed for Zstandard decompression using a
2228 custom compression dictionary:
2229 ``compression={'method': 'zstd', 'dict_data': my_compression_dict}``.
2230 storage_options : dict, optional
2231 Extra options that make sense for a particular storage connection, e.g.
2232 host, port, username, password, etc. For HTTP(S) URLs the key-value pairs
2233 are forwarded to ``urllib.request.Request`` as header options. For other
2234 URLs (e.g. starting with "s3://", and "gcs://") the key-value pairs are
2235 forwarded to ``fsspec.open``. Please see ``fsspec`` and ``urllib`` for more
2236 details, and for more examples on storage options refer `here
2237 <https://pandas.pydata.org/docs/user_guide/io.html?
2238 highlight=storage_options#reading-writing-remote-files>`_.
2240 Returns
2241 -------
2242 DataFrame, pandas.api.typing.StataReader
2243 If iterator or chunksize, returns StataReader, else DataFrame.
2245 See Also
2246 --------
2247 io.stata.StataReader : Low-level reader for Stata data files.
2248 DataFrame.to_stata: Export Stata data files.
2250 Notes
2251 -----
2252 Categorical variables read through an iterator may not have the same
2253 categories and dtype. This occurs when a variable stored in a DTA
2254 file is associated to an incomplete set of value labels that only
2255 label a strict subset of the values.
2257 Examples
2258 --------
2260 Creating a dummy stata for this example
2262 >>> df = pd.DataFrame(
2263 ... {
2264 ... "animal": ["falcon", "parrot", "falcon", "parrot"],
2265 ... "speed": [350, 18, 361, 15],
2266 ... }
2267 ... ) # doctest: +SKIP
2268 >>> df.to_stata("animals.dta") # doctest: +SKIP
2270 Read a Stata dta file:
2272 >>> df = pd.read_stata("animals.dta") # doctest: +SKIP
2274 Read a Stata dta file in 10,000 line chunks:
2276 >>> values = np.random.randint(
2277 ... 0, 10, size=(20_000, 1), dtype="uint8"
2278 ... ) # doctest: +SKIP
2279 >>> df = pd.DataFrame(values, columns=["i"]) # doctest: +SKIP
2280 >>> df.to_stata("filename.dta") # doctest: +SKIP
2282 >>> with pd.read_stata('filename.dta', chunksize=10000) as itr: # doctest: +SKIP
2283 >>> for chunk in itr:
2284 ... # Operate on a single chunk, e.g., chunk.mean()
2285 ... pass # doctest: +SKIP
2286 """
2287 reader = StataReader(
2288 filepath_or_buffer,
2289 convert_dates=convert_dates,
2290 convert_categoricals=convert_categoricals,
2291 index_col=index_col,
2292 convert_missing=convert_missing,
2293 preserve_dtypes=preserve_dtypes,
2294 columns=columns,
2295 order_categoricals=order_categoricals,
2296 chunksize=chunksize,
2297 storage_options=storage_options,
2298 compression=compression,
2299 )
2301 if iterator or chunksize:
2302 return reader
2304 with reader:
2305 return reader.read()
2308def _set_endianness(endianness: str) -> str:
2309 if endianness.lower() in ["<", "little"]:
2310 return "<"
2311 elif endianness.lower() in [">", "big"]:
2312 return ">"
2313 else: # pragma : no cover
2314 raise ValueError(f"Endianness {endianness} not understood")
2317def _pad_bytes(name: AnyStr, length: int) -> AnyStr:
2318 """
2319 Take a char string and pads it with null bytes until it's length chars.
2320 """
2321 if isinstance(name, bytes):
2322 return name + b"\x00" * (length - len(name))
2323 return name + "\x00" * (length - len(name))
2326def _convert_datetime_to_stata_type(fmt: str) -> np.dtype:
2327 """
2328 Convert from one of the stata date formats to a type in TYPE_MAP.
2329 """
2330 if fmt in [
2331 "tc",
2332 "%tc",
2333 "td",
2334 "%td",
2335 "tw",
2336 "%tw",
2337 "tm",
2338 "%tm",
2339 "tq",
2340 "%tq",
2341 "th",
2342 "%th",
2343 "ty",
2344 "%ty",
2345 ]:
2346 return np.dtype(np.float64) # Stata expects doubles for SIFs
2347 else:
2348 raise NotImplementedError(f"Format {fmt} not implemented")
2351def _maybe_convert_to_int_keys(
2352 convert_dates: dict[Hashable, str], varlist: list[Hashable]
2353) -> dict[Hashable, str]:
2354 new_dict: dict[Hashable, str] = {}
2355 for key, value in convert_dates.items():
2356 if not value.startswith("%"): # make sure proper fmts
2357 convert_dates[key] = "%" + value
2358 if key in varlist:
2359 new_dict[varlist.index(key)] = convert_dates[key]
2360 else:
2361 if not isinstance(key, int):
2362 raise ValueError("convert_dates key must be a column or an integer")
2363 new_dict[key] = convert_dates[key]
2364 return new_dict
2367def _dtype_to_stata_type(dtype: np.dtype, column: Series) -> int:
2368 """
2369 Convert dtype types to stata types. Returns the byte of the given ordinal.
2370 See TYPE_MAP and comments for an explanation. This is also explained in
2371 the dta spec.
2372 1 - 244 are strings of this length
2373 Pandas Stata
2374 251 - for int8 byte
2375 252 - for int16 int
2376 253 - for int32 long
2377 254 - for float32 float
2378 255 - for double double
2380 If there are dates to convert, then dtype will already have the correct
2381 type inserted.
2382 """
2383 # TODO: expand to handle datetime to integer conversion
2384 if dtype.type is np.object_: # try to coerce it to the biggest string
2385 # not memory efficient, what else could we
2386 # do?
2387 itemsize = max_len_string_array(ensure_object(column._values))
2388 return max(itemsize, 1)
2389 elif dtype.type is np.float64:
2390 return 255
2391 elif dtype.type is np.float32:
2392 return 254
2393 elif dtype.type is np.int32:
2394 return 253
2395 elif dtype.type is np.int16:
2396 return 252
2397 elif dtype.type is np.int8:
2398 return 251
2399 else: # pragma : no cover
2400 raise NotImplementedError(f"Data type {dtype} not supported.")
2403def _dtype_to_default_stata_fmt(
2404 dtype: np.dtype, column: Series, dta_version: int = 114, force_strl: bool = False
2405) -> str:
2406 """
2407 Map numpy dtype to stata's default format for this type. Not terribly
2408 important since users can change this in Stata. Semantics are
2410 object -> "%DDs" where DD is the length of the string. If not a string,
2411 raise ValueError
2412 float64 -> "%10.0g"
2413 float32 -> "%9.0g"
2414 int64 -> "%9.0g"
2415 int32 -> "%12.0g"
2416 int16 -> "%8.0g"
2417 int8 -> "%8.0g"
2418 strl -> "%9s"
2419 """
2420 # TODO: Refactor to combine type with format
2421 # TODO: expand this to handle a default datetime format?
2422 if dta_version < 117:
2423 max_str_len = 244
2424 else:
2425 max_str_len = 2045
2426 if force_strl:
2427 return "%9s"
2428 if dtype.type is np.object_:
2429 itemsize = max_len_string_array(ensure_object(column._values))
2430 if itemsize > max_str_len:
2431 if dta_version >= 117:
2432 return "%9s"
2433 else:
2434 raise ValueError(excessive_string_length_error.format(column.name))
2435 return "%" + str(max(itemsize, 1)) + "s"
2436 elif dtype == np.float64:
2437 return "%10.0g"
2438 elif dtype == np.float32:
2439 return "%9.0g"
2440 elif dtype == np.int32:
2441 return "%12.0g"
2442 elif dtype in (np.int8, np.int16):
2443 return "%8.0g"
2444 else: # pragma : no cover
2445 raise NotImplementedError(f"Data type {dtype} not supported.")
2448class StataWriter(StataParser):
2449 """
2450 A class for writing Stata binary dta files
2452 Parameters
2453 ----------
2454 fname : path (string), buffer or path object
2455 string, pathlib.Path or
2456 object implementing a binary write() functions. If using a buffer
2457 then the buffer will not be automatically closed after the file
2458 is written.
2459 data : DataFrame
2460 Input to save
2461 convert_dates : dict
2462 Dictionary mapping columns containing datetime types to stata internal
2463 format to use when writing the dates. Options are 'tc', 'td', 'tm',
2464 'tw', 'th', 'tq', 'ty'. Column can be either an integer or a name.
2465 Datetime columns that do not have a conversion type specified will be
2466 converted to 'tc'. Raises NotImplementedError if a datetime column has
2467 timezone information
2468 write_index : bool
2469 Write the index to Stata dataset.
2470 byteorder : str
2471 Can be ">", "<", "little", or "big". default is `sys.byteorder`
2472 time_stamp : datetime
2473 A datetime to use as file creation date. Default is the current time
2474 data_label : str
2475 A label for the data set. Must be 80 characters or smaller.
2476 variable_labels : dict
2477 Dictionary containing columns as keys and variable labels as values.
2478 Each label must be 80 characters or smaller.
2479 compression : str or dict, default 'infer'
2480 For on-the-fly compression of the output data. If 'infer' and 'fname' is
2481 path-like, then detect compression from the following extensions: '.gz',
2482 '.bz2', '.zip', '.xz', '.zst', '.tar', '.tar.gz', '.tar.xz' or '.tar.bz2'
2483 (otherwise no compression).
2484 Set to ``None`` for no compression.
2485 Can also be a dict with key ``'method'`` set
2486 to one of {``'zip'``, ``'gzip'``, ``'bz2'``, ``'zstd'``, ``'xz'``, ``'tar'``}
2487 and other key-value pairs are forwarded to
2488 ``zipfile.ZipFile``, ``gzip.GzipFile``,
2489 ``bz2.BZ2File``, ``zstandard.ZstdCompressor``, ``lzma.LZMAFile`` or
2490 ``tarfile.TarFile``, respectively.
2491 As an example, the following could be passed for faster compression and to
2492 create a reproducible gzip archive:
2493 ``compression={'method': 'gzip', 'compresslevel': 1, 'mtime': 1}``.
2494 storage_options : dict, optional
2495 Extra options that make sense for a particular storage connection, e.g.
2496 host, port, username, password, etc. For HTTP(S) URLs the key-value pairs
2497 are forwarded to ``urllib.request.Request`` as header options. For other
2498 URLs (e.g. starting with "s3://", and "gcs://") the key-value pairs are
2499 forwarded to ``fsspec.open``. Please see ``fsspec`` and ``urllib`` for more
2500 details, and for more examples on storage options refer `here
2501 <https://pandas.pydata.org/docs/user_guide/io.html?
2502 highlight=storage_options#reading-writing-remote-files>`_.
2504 value_labels : dict of dicts
2505 Dictionary containing columns as keys and dictionaries of column value
2506 to labels as values. The combined length of all labels for a single
2507 variable must be 32,000 characters or smaller.
2509 Returns
2510 -------
2511 writer : StataWriter instance
2512 The StataWriter instance has a write_file method, which will
2513 write the file to the given `fname`.
2515 Raises
2516 ------
2517 NotImplementedError
2518 * If datetimes contain timezone information
2519 ValueError
2520 * Columns listed in convert_dates are neither datetime64[ns]
2521 or datetime
2522 * Column dtype is not representable in Stata
2523 * Column listed in convert_dates is not in DataFrame
2524 * Categorical label contains more than 32,000 characters
2526 Examples
2527 --------
2528 >>> data = pd.DataFrame([[1.0, 1]], columns=["a", "b"])
2529 >>> writer = StataWriter("./data_file.dta", data)
2530 >>> writer.write_file()
2532 Directly write a zip file
2533 >>> compression = {"method": "zip", "archive_name": "data_file.dta"}
2534 >>> writer = StataWriter("./data_file.zip", data, compression=compression)
2535 >>> writer.write_file()
2537 Save a DataFrame with dates
2538 >>> from datetime import datetime
2539 >>> data = pd.DataFrame([[datetime(2000, 1, 1)]], columns=["date"])
2540 >>> writer = StataWriter("./date_data_file.dta", data, {"date": "tw"})
2541 >>> writer.write_file()
2542 """
2544 _max_string_length = 244
2545 _encoding: Literal["latin-1", "utf-8"] = "latin-1"
2547 def __init__(
2548 self,
2549 fname: FilePath | WriteBuffer[bytes],
2550 data: DataFrame,
2551 convert_dates: dict[Hashable, str] | None = None,
2552 write_index: bool = True,
2553 byteorder: str | None = None,
2554 time_stamp: datetime | None = None,
2555 data_label: str | None = None,
2556 variable_labels: dict[Hashable, str] | None = None,
2557 compression: CompressionOptions = "infer",
2558 storage_options: StorageOptions | None = None,
2559 *,
2560 value_labels: dict[Hashable, dict[float, str]] | None = None,
2561 ) -> None:
2562 super().__init__()
2563 self.data = data
2564 self._convert_dates = {} if convert_dates is None else convert_dates
2565 self._write_index = write_index
2566 self._time_stamp = time_stamp
2567 self._data_label = data_label
2568 self._variable_labels = variable_labels
2569 self._non_cat_value_labels = value_labels
2570 self._value_labels: list[StataValueLabel] = []
2571 self._has_value_labels = np.array([], dtype=bool)
2572 self._compression = compression
2573 self._output_file: IO[bytes] | None = None
2574 self._converted_names: dict[Hashable, str] = {}
2575 # attach nobs, nvars, data, varlist, typlist
2576 self._prepare_pandas(data)
2577 self.storage_options = storage_options
2579 if byteorder is None:
2580 byteorder = sys.byteorder
2581 self._byteorder = _set_endianness(byteorder)
2582 self._fname = fname
2583 self.type_converters = {253: np.int32, 252: np.int16, 251: np.int8}
2585 def _write(self, to_write: str) -> None:
2586 """
2587 Helper to call encode before writing to file for Python 3 compat.
2588 """
2589 self.handles.handle.write(to_write.encode(self._encoding))
2591 def _write_bytes(self, value: bytes) -> None:
2592 """
2593 Helper to assert file is open before writing.
2594 """
2595 self.handles.handle.write(value)
2597 def _prepare_non_cat_value_labels(
2598 self, data: DataFrame
2599 ) -> list[StataNonCatValueLabel]:
2600 """
2601 Check for value labels provided for non-categorical columns. Value
2602 labels
2603 """
2604 non_cat_value_labels: list[StataNonCatValueLabel] = []
2605 if self._non_cat_value_labels is None:
2606 return non_cat_value_labels
2608 for labname, labels in self._non_cat_value_labels.items():
2609 if labname in self._converted_names:
2610 colname = self._converted_names[labname]
2611 elif labname in data.columns:
2612 colname = str(labname)
2613 else:
2614 raise KeyError(
2615 f"Can't create value labels for {labname}, it wasn't "
2616 "found in the dataset."
2617 )
2619 if not is_numeric_dtype(data[colname].dtype):
2620 # Labels should not be passed explicitly for categorical
2621 # columns that will be converted to int
2622 raise ValueError(
2623 f"Can't create value labels for {labname}, value labels "
2624 "can only be applied to numeric columns."
2625 )
2626 svl = StataNonCatValueLabel(colname, labels, self._encoding)
2627 non_cat_value_labels.append(svl)
2628 return non_cat_value_labels
2630 def _prepare_categoricals(self, data: DataFrame) -> DataFrame:
2631 """
2632 Check for categorical columns, retain categorical information for
2633 Stata file and convert categorical data to int
2634 """
2635 is_cat = [isinstance(dtype, CategoricalDtype) for dtype in data.dtypes]
2636 if not any(is_cat):
2637 return data
2639 self._has_value_labels |= np.array(is_cat)
2641 get_base_missing_value = StataMissingValue.get_base_missing_value
2642 data_formatted = []
2643 for col, col_is_cat in zip(data, is_cat, strict=True):
2644 if col_is_cat:
2645 svl = StataValueLabel(data[col], encoding=self._encoding)
2646 self._value_labels.append(svl)
2647 dtype = data[col].cat.codes.dtype
2648 if dtype == np.int64:
2649 raise ValueError(
2650 "It is not possible to export "
2651 "int64-based categorical data to Stata."
2652 )
2653 values = data[col].cat.codes._values.copy()
2655 # Upcast if needed so that correct missing values can be set
2656 if values.max() >= get_base_missing_value(dtype):
2657 if dtype == np.int8:
2658 dtype = np.dtype(np.int16)
2659 elif dtype == np.int16:
2660 dtype = np.dtype(np.int32)
2661 else:
2662 dtype = np.dtype(np.float64)
2663 values = np.array(values, dtype=dtype)
2665 # Replace missing values with Stata missing value for type
2666 values[values == -1] = get_base_missing_value(dtype)
2667 data_formatted.append((col, values))
2668 else:
2669 data_formatted.append((col, data[col]))
2670 return DataFrame.from_dict(dict(data_formatted))
2672 def _replace_nans(self, data: DataFrame) -> DataFrame:
2673 # return data
2674 """
2675 Checks floating point data columns for nans, and replaces these with
2676 the generic Stata for missing value (.)
2677 """
2678 for c in data:
2679 dtype = data[c].dtype
2680 if dtype in (np.float32, np.float64):
2681 if dtype == np.float32:
2682 replacement = self.MISSING_VALUES["f"]
2683 else:
2684 replacement = self.MISSING_VALUES["d"]
2685 data[c] = data[c].fillna(replacement)
2687 return data
2689 def _update_strl_names(self) -> None:
2690 """No-op, forward compatibility"""
2692 def _validate_variable_name(self, name: str) -> str:
2693 """
2694 Validate variable names for Stata export.
2696 Parameters
2697 ----------
2698 name : str
2699 Variable name
2701 Returns
2702 -------
2703 str
2704 The validated name with invalid characters replaced with
2705 underscores.
2707 Notes
2708 -----
2709 Stata 114 and 117 support ascii characters in a-z, A-Z, 0-9
2710 and _.
2711 """
2712 for c in name:
2713 if (
2714 (c < "A" or c > "Z")
2715 and (c < "a" or c > "z")
2716 and (c < "0" or c > "9")
2717 and c != "_"
2718 ):
2719 name = name.replace(c, "_")
2720 return name
2722 def _check_column_names(self, data: DataFrame) -> DataFrame:
2723 """
2724 Checks column names to ensure that they are valid Stata column names.
2725 This includes checks for:
2726 * Non-string names
2727 * Stata keywords
2728 * Variables that start with numbers
2729 * Variables with names that are too long
2731 When an illegal variable name is detected, it is converted, and if
2732 dates are exported, the variable name is propagated to the date
2733 conversion dictionary
2734 """
2735 converted_names: dict[Hashable, str] = {}
2736 columns = list(data.columns)
2737 original_columns = columns[:]
2739 duplicate_var_id = 0
2740 for j, name in enumerate(columns):
2741 orig_name = name
2742 if not isinstance(name, str):
2743 name = str(name)
2745 name = self._validate_variable_name(name)
2747 # Variable name must not be a reserved word
2748 if name in self.RESERVED_WORDS:
2749 name = "_" + name
2751 # Variable name may not start with a number
2752 if "0" <= name[0] <= "9":
2753 name = "_" + name
2755 name = name[: min(len(name), 32)]
2757 if not name == orig_name:
2758 # check for duplicates
2759 while columns.count(name) > 0:
2760 # prepend ascending number to avoid duplicates
2761 name = "_" + str(duplicate_var_id) + name
2762 name = name[: min(len(name), 32)]
2763 duplicate_var_id += 1
2764 converted_names[orig_name] = name
2766 columns[j] = name
2768 data.columns = Index(columns)
2770 # Check date conversion, and fix key if needed
2771 if self._convert_dates:
2772 for c, o in zip(columns, original_columns, strict=True):
2773 if c != o and o in self._convert_dates:
2774 self._convert_dates[c] = self._convert_dates[o]
2775 del self._convert_dates[o]
2777 if converted_names:
2778 conversion_warning = []
2779 for orig_name, name in converted_names.items():
2780 msg = f"{orig_name} -> {name}"
2781 conversion_warning.append(msg)
2783 ws = invalid_name_doc.format("\n ".join(conversion_warning))
2784 warnings.warn(
2785 ws,
2786 InvalidColumnName,
2787 stacklevel=find_stack_level(),
2788 )
2790 self._converted_names = converted_names
2791 self._update_strl_names()
2793 return data
2795 def _set_formats_and_types(self, dtypes: Series) -> None:
2796 self.fmtlist: list[str] = []
2797 self.typlist: list[int] = []
2798 for col, dtype in dtypes.items():
2799 self.fmtlist.append(_dtype_to_default_stata_fmt(dtype, self.data[col]))
2800 self.typlist.append(_dtype_to_stata_type(dtype, self.data[col]))
2802 def _prepare_pandas(self, data: DataFrame) -> None:
2803 # NOTE: we might need a different API / class for pandas objects so
2804 # we can set different semantics - handle this with a PR to pandas.io
2806 data = data.copy()
2808 if self._write_index:
2809 temp = data.reset_index()
2810 if isinstance(temp, DataFrame):
2811 data = temp
2813 # Ensure column names are strings
2814 data = self._check_column_names(data)
2816 # Check columns for compatibility with stata, upcast if necessary
2817 # Raise if outside the supported range
2818 data = _cast_to_stata_types(data)
2820 # Replace NaNs with Stata missing values
2821 data = self._replace_nans(data)
2823 # Set all columns to initially unlabelled
2824 self._has_value_labels = np.repeat(False, data.shape[1])
2826 # Create value labels for non-categorical data
2827 non_cat_value_labels = self._prepare_non_cat_value_labels(data)
2829 non_cat_columns = [svl.labname for svl in non_cat_value_labels]
2830 has_non_cat_val_labels = data.columns.isin(non_cat_columns)
2831 self._has_value_labels |= has_non_cat_val_labels
2832 self._value_labels.extend(non_cat_value_labels)
2834 # Convert categoricals to int data, and strip labels
2835 data = self._prepare_categoricals(data)
2837 self.nobs, self.nvar = data.shape
2838 self.data = data
2839 self.varlist = data.columns.tolist()
2841 dtypes = data.dtypes
2843 # Ensure all date columns are converted
2844 for col in data:
2845 if col in self._convert_dates:
2846 continue
2847 if lib.is_np_dtype(data[col].dtype, "M"):
2848 self._convert_dates[col] = "tc"
2850 self._convert_dates = _maybe_convert_to_int_keys(
2851 self._convert_dates, self.varlist
2852 )
2853 for key in self._convert_dates:
2854 new_type = _convert_datetime_to_stata_type(self._convert_dates[key])
2855 dtypes.iloc[key] = np.dtype(new_type)
2857 # Verify object arrays are strings and encode to bytes
2858 self._encode_strings()
2860 self._set_formats_and_types(dtypes)
2862 # set the given format for the datetime cols
2863 if self._convert_dates is not None:
2864 for key in self._convert_dates:
2865 if isinstance(key, int):
2866 self.fmtlist[key] = self._convert_dates[key]
2868 def _encode_strings(self) -> None:
2869 """
2870 Encode strings in dta-specific encoding
2872 Do not encode columns marked for date conversion or for strL
2873 conversion. The strL converter independently handles conversion and
2874 also accepts empty string arrays.
2875 """
2876 convert_dates = self._convert_dates
2877 # _convert_strl is not available in dta 114
2878 convert_strl = getattr(self, "_convert_strl", [])
2879 for i, col in enumerate(self.data):
2880 # Skip columns marked for date conversion or strl conversion
2881 if i in convert_dates or col in convert_strl:
2882 continue
2883 column = self.data[col]
2884 dtype = column.dtype
2885 # TODO could also handle string dtype here specifically
2886 if dtype.type is np.object_:
2887 inferred_dtype = infer_dtype(column, skipna=True)
2888 if not ((inferred_dtype == "string") or len(column) == 0):
2889 col = column.name
2890 raise ValueError(
2891 f"""\
2892Column `{col}` cannot be exported.\n\nOnly string-like object arrays
2893containing all strings or a mix of strings and None can be exported.
2894Object arrays containing only null values are prohibited. Other object
2895types cannot be exported and must first be converted to one of the
2896supported types."""
2897 )
2898 encoded = self.data[col].str.encode(self._encoding)
2899 # If larger than _max_string_length do nothing
2900 if (
2901 max_len_string_array(ensure_object(self.data[col]._values))
2902 <= self._max_string_length
2903 ):
2904 self.data[col] = encoded
2906 def write_file(self) -> None:
2907 """
2908 Export DataFrame object to Stata dta format.
2910 This method writes the contents of a pandas DataFrame to a `.dta` file
2911 compatible with Stata. It includes features for handling value labels,
2912 variable types, and metadata like timestamps and data labels. The output
2913 file can then be read and used in Stata or other compatible statistical
2914 tools.
2916 See Also
2917 --------
2918 read_stata : Read Stata file into DataFrame.
2919 DataFrame.to_stata : Export DataFrame object to Stata dta format.
2920 io.stata.StataWriter : A class for writing Stata binary dta files.
2922 Examples
2923 --------
2924 >>> df = pd.DataFrame(
2925 ... {
2926 ... "fully_labelled": [1, 2, 3, 3, 1],
2927 ... "partially_labelled": [1.0, 2.0, np.nan, 9.0, np.nan],
2928 ... "Y": [7, 7, 9, 8, 10],
2929 ... "Z": pd.Categorical(["j", "k", "l", "k", "j"]),
2930 ... }
2931 ... )
2932 >>> path = "/My_path/filename.dta"
2933 >>> labels = {
2934 ... "fully_labelled": {1: "one", 2: "two", 3: "three"},
2935 ... "partially_labelled": {1.0: "one", 2.0: "two"},
2936 ... }
2937 >>> writer = pd.io.stata.StataWriter(
2938 ... path, df, value_labels=labels
2939 ... ) # doctest: +SKIP
2940 >>> writer.write_file() # doctest: +SKIP
2941 >>> df = pd.read_stata(path) # doctest: +SKIP
2942 >>> df # doctest: +SKIP
2943 index fully_labelled partially_labeled Y Z
2944 0 0 one one 7 j
2945 1 1 two two 7 k
2946 2 2 three NaN 9 l
2947 3 3 three 9.0 8 k
2948 4 4 one NaN 10 j
2949 """
2950 with get_handle(
2951 self._fname,
2952 "wb",
2953 compression=self._compression,
2954 is_text=False,
2955 storage_options=self.storage_options,
2956 ) as self.handles:
2957 if self.handles.compression["method"] is not None:
2958 # ZipFile creates a file (with the same name) for each write call.
2959 # Write it first into a buffer and then write the buffer to the ZipFile.
2960 self._output_file, self.handles.handle = self.handles.handle, BytesIO()
2961 self.handles.created_handles.append(self.handles.handle)
2963 try:
2964 self._write_header(
2965 data_label=self._data_label, time_stamp=self._time_stamp
2966 )
2967 self._write_map()
2968 self._write_variable_types()
2969 self._write_varnames()
2970 self._write_sortlist()
2971 self._write_formats()
2972 self._write_value_label_names()
2973 self._write_variable_labels()
2974 self._write_expansion_fields()
2975 self._write_characteristics()
2976 records = self._prepare_data()
2977 self._write_data(records)
2978 self._write_strls()
2979 self._write_value_labels()
2980 self._write_file_close_tag()
2981 self._write_map()
2982 self._close()
2983 except Exception as exc:
2984 self.handles.close()
2985 if isinstance(self._fname, (str, os.PathLike)) and os.path.isfile(
2986 self._fname
2987 ):
2988 try:
2989 os.unlink(self._fname)
2990 except OSError:
2991 warnings.warn(
2992 f"This save was not successful but {self._fname} could not "
2993 "be deleted. This file is not valid.",
2994 ResourceWarning,
2995 stacklevel=find_stack_level(),
2996 )
2997 raise exc
2999 def _close(self) -> None:
3000 """
3001 Close the file if it was created by the writer.
3003 If a buffer or file-like object was passed in, for example a GzipFile,
3004 then leave this file open for the caller to close.
3005 """
3006 # write compression
3007 if self._output_file is not None:
3008 assert isinstance(self.handles.handle, BytesIO)
3009 bio, self.handles.handle = self.handles.handle, self._output_file
3010 self.handles.handle.write(bio.getvalue())
3012 def _write_map(self) -> None:
3013 """No-op, future compatibility"""
3015 def _write_file_close_tag(self) -> None:
3016 """No-op, future compatibility"""
3018 def _write_characteristics(self) -> None:
3019 """No-op, future compatibility"""
3021 def _write_strls(self) -> None:
3022 """No-op, future compatibility"""
3024 def _write_expansion_fields(self) -> None:
3025 """Write 5 zeros for expansion fields"""
3026 self._write(_pad_bytes("", 5))
3028 def _write_value_labels(self) -> None:
3029 for vl in self._value_labels:
3030 self._write_bytes(vl.generate_value_label(self._byteorder))
3032 def _write_header(
3033 self,
3034 data_label: str | None = None,
3035 time_stamp: datetime | None = None,
3036 ) -> None:
3037 byteorder = self._byteorder
3038 # ds_format - just use 114
3039 self._write_bytes(struct.pack("b", 114))
3040 # byteorder
3041 self._write((byteorder == ">" and "\x01") or "\x02")
3042 # filetype
3043 self._write("\x01")
3044 # unused
3045 self._write("\x00")
3046 # number of vars, 2 bytes
3047 self._write_bytes(struct.pack(byteorder + "h", self.nvar)[:2])
3048 # number of obs, 4 bytes
3049 self._write_bytes(struct.pack(byteorder + "i", self.nobs)[:4])
3050 # data label 81 bytes, char, null terminated
3051 if data_label is None:
3052 self._write_bytes(self._null_terminate_bytes(_pad_bytes("", 80)))
3053 else:
3054 self._write_bytes(
3055 self._null_terminate_bytes(_pad_bytes(data_label[:80], 80))
3056 )
3057 # time stamp, 18 bytes, char, null terminated
3058 # format dd Mon yyyy hh:mm
3059 if time_stamp is None:
3060 time_stamp = datetime.now() # noqa: TID251
3061 elif not isinstance(time_stamp, datetime):
3062 raise ValueError("time_stamp should be datetime type")
3063 # GH #13856
3064 # Avoid locale-specific month conversion
3065 months = [
3066 "Jan",
3067 "Feb",
3068 "Mar",
3069 "Apr",
3070 "May",
3071 "Jun",
3072 "Jul",
3073 "Aug",
3074 "Sep",
3075 "Oct",
3076 "Nov",
3077 "Dec",
3078 ]
3079 month_lookup = {i + 1: month for i, month in enumerate(months)}
3080 ts = (
3081 time_stamp.strftime("%d ")
3082 + month_lookup[time_stamp.month]
3083 + time_stamp.strftime(" %Y %H:%M")
3084 )
3085 self._write_bytes(self._null_terminate_bytes(ts))
3087 def _write_variable_types(self) -> None:
3088 for typ in self.typlist:
3089 self._write_bytes(struct.pack("B", typ))
3091 def _write_varnames(self) -> None:
3092 # varlist names are checked by _check_column_names
3093 # varlist, requires null terminated
3094 for name in self.varlist:
3095 name = self._null_terminate_str(name)
3096 name = _pad_bytes(name[:32], 33)
3097 self._write(name)
3099 def _write_sortlist(self) -> None:
3100 # srtlist, 2*(nvar+1), int array, encoded by byteorder
3101 srtlist = _pad_bytes("", 2 * (self.nvar + 1))
3102 self._write(srtlist)
3104 def _write_formats(self) -> None:
3105 # fmtlist, 49*nvar, char array
3106 for fmt in self.fmtlist:
3107 self._write(_pad_bytes(fmt, 49))
3109 def _write_value_label_names(self) -> None:
3110 # lbllist, 33*nvar, char array
3111 for i in range(self.nvar):
3112 # Use variable name when categorical
3113 if self._has_value_labels[i]:
3114 name = self.varlist[i]
3115 name = self._null_terminate_str(name)
3116 name = _pad_bytes(name[:32], 33)
3117 self._write(name)
3118 else: # Default is empty label
3119 self._write(_pad_bytes("", 33))
3121 def _write_variable_labels(self) -> None:
3122 # Missing labels are 80 blank characters plus null termination
3123 blank = _pad_bytes("", 81)
3125 if self._variable_labels is None:
3126 for i in range(self.nvar):
3127 self._write(blank)
3128 return
3130 for col in self.data:
3131 if col in self._variable_labels:
3132 label = self._variable_labels[col]
3133 if len(label) > 80:
3134 raise ValueError("Variable labels must be 80 characters or fewer")
3135 is_latin1 = all(ord(c) < 256 for c in label)
3136 if not is_latin1:
3137 raise ValueError(
3138 "Variable labels must contain only characters that "
3139 "can be encoded in Latin-1"
3140 )
3141 self._write(_pad_bytes(label, 81))
3142 else:
3143 self._write(blank)
3145 def _convert_strls(self, data: DataFrame) -> DataFrame:
3146 """No-op, future compatibility"""
3147 return data
3149 def _prepare_data(self) -> np.rec.recarray:
3150 data = self.data
3151 typlist = self.typlist
3152 convert_dates = self._convert_dates
3153 # 1. Convert dates
3154 if self._convert_dates is not None:
3155 for i, col in enumerate(data):
3156 if i in convert_dates:
3157 data[col] = _datetime_to_stata_elapsed_vec(
3158 data[col], self.fmtlist[i]
3159 )
3160 # 2. Convert strls
3161 data = self._convert_strls(data)
3163 # 3. Convert bad string data to '' and pad to correct length
3164 dtypes = {}
3165 native_byteorder = self._byteorder == _set_endianness(sys.byteorder)
3166 for i, col in enumerate(data):
3167 typ = typlist[i]
3168 if typ <= self._max_string_length:
3169 dc = data[col].fillna("")
3170 data[col] = dc.apply(_pad_bytes, args=(typ,))
3171 stype = f"S{typ}"
3172 dtypes[col] = stype
3173 data[col] = data[col].astype(stype)
3174 else:
3175 dtype = data[col].dtype
3176 if not native_byteorder:
3177 dtype = dtype.newbyteorder(self._byteorder)
3178 dtypes[col] = dtype
3180 return data.to_records(index=False, column_dtypes=dtypes)
3182 def _write_data(self, records: np.rec.recarray) -> None:
3183 self._write_bytes(records.tobytes())
3185 @staticmethod
3186 def _null_terminate_str(s: str) -> str:
3187 s += "\x00"
3188 return s
3190 def _null_terminate_bytes(self, s: str) -> bytes:
3191 return self._null_terminate_str(s).encode(self._encoding)
3194def _dtype_to_stata_type_117(dtype: np.dtype, column: Series, force_strl: bool) -> int:
3195 """
3196 Converts dtype types to stata types. Returns the byte of the given ordinal.
3197 See TYPE_MAP and comments for an explanation. This is also explained in
3198 the dta spec.
3199 1 - 2045 are strings of this length
3200 Pandas Stata
3201 32768 - for object strL
3202 65526 - for int8 byte
3203 65527 - for int16 int
3204 65528 - for int32 long
3205 65529 - for float32 float
3206 65530 - for double double
3208 If there are dates to convert, then dtype will already have the correct
3209 type inserted.
3210 """
3211 # TODO: expand to handle datetime to integer conversion
3212 if force_strl:
3213 return 32768
3214 if dtype.type is np.object_: # try to coerce it to the biggest string
3215 # not memory efficient, what else could we
3216 # do?
3217 itemsize = max_len_string_array(ensure_object(column._values))
3218 itemsize = max(itemsize, 1)
3219 if itemsize <= 2045:
3220 return itemsize
3221 return 32768
3222 elif dtype.type is np.float64:
3223 return 65526
3224 elif dtype.type is np.float32:
3225 return 65527
3226 elif dtype.type is np.int32:
3227 return 65528
3228 elif dtype.type is np.int16:
3229 return 65529
3230 elif dtype.type is np.int8:
3231 return 65530
3232 else: # pragma : no cover
3233 raise NotImplementedError(f"Data type {dtype} not supported.")
3236def _pad_bytes_new(name: str | bytes, length: int) -> bytes:
3237 """
3238 Takes a bytes instance and pads it with null bytes until it's length chars.
3239 """
3240 if isinstance(name, str):
3241 name = bytes(name, "utf-8")
3242 return name + b"\x00" * (length - len(name))
3245class StataStrLWriter:
3246 """
3247 Converter for Stata StrLs
3249 Stata StrLs map 8 byte values to strings which are stored using a
3250 dictionary-like format where strings are keyed to two values.
3252 Parameters
3253 ----------
3254 df : DataFrame
3255 DataFrame to convert
3256 columns : Sequence[str]
3257 List of columns names to convert to StrL
3258 version : int, optional
3259 dta version. Currently supports 117, 118 and 119
3260 byteorder : str, optional
3261 Can be ">", "<", "little", or "big". default is `sys.byteorder`
3263 Notes
3264 -----
3265 Supports creation of the StrL block of a dta file for dta versions
3266 117, 118 and 119. These differ in how the GSO is stored. 118 and
3267 119 store the GSO lookup value as a uint32 and a uint64, while 117
3268 uses two uint32s. 118 and 119 also encode all strings as unicode
3269 which is required by the format. 117 uses 'latin-1' a fixed width
3270 encoding that extends the 7-bit ascii table with an additional 128
3271 characters.
3272 """
3274 def __init__(
3275 self,
3276 df: DataFrame,
3277 columns: Sequence[str],
3278 version: int = 117,
3279 byteorder: str | None = None,
3280 ) -> None:
3281 if version not in (117, 118, 119):
3282 raise ValueError("Only dta versions 117, 118 and 119 supported")
3283 self._dta_ver = version
3285 self.df = df
3286 self.columns = columns
3287 self._gso_table = {"": (0, 0)}
3288 if byteorder is None:
3289 byteorder = sys.byteorder
3290 self._byteorder = _set_endianness(byteorder)
3291 # Flag whether chosen byteorder matches the system on which we're running
3292 self._native_byteorder = self._byteorder == _set_endianness(sys.byteorder)
3294 gso_v_type = "I" # uint32
3295 gso_o_type = "Q" # uint64
3296 self._encoding = "utf-8"
3297 if version == 117:
3298 o_size = 4
3299 gso_o_type = "I" # 117 used uint32
3300 self._encoding = "latin-1"
3301 elif version == 118:
3302 o_size = 6
3303 else: # version == 119
3304 o_size = 5
3305 if self._native_byteorder:
3306 self._o_offet = 2 ** (8 * (8 - o_size))
3307 else:
3308 self._o_offet = 2 ** (8 * o_size)
3309 self._gso_o_type = gso_o_type
3310 self._gso_v_type = gso_v_type
3312 def _convert_key(self, key: tuple[int, int]) -> int:
3313 v, o = key
3314 if self._native_byteorder:
3315 return v + self._o_offet * o
3316 else:
3317 # v, o will be swapped when applying byteorder
3318 return o + self._o_offet * v
3320 def generate_table(self) -> tuple[dict[str, tuple[int, int]], DataFrame]:
3321 """
3322 Generates the GSO lookup table for the DataFrame
3324 Returns
3325 -------
3326 gso_table : dict
3327 Ordered dictionary using the string found as keys
3328 and their lookup position (v,o) as values
3329 gso_df : DataFrame
3330 DataFrame where strl columns have been converted to
3331 (v,o) values
3333 Notes
3334 -----
3335 Modifies the DataFrame in-place.
3337 The DataFrame returned encodes the (v,o) values as uint64s. The
3338 encoding depends on the dta version, and can be expressed as
3340 enc = v + o * 2 ** (o_size * 8)
3342 so that v is stored in the lower bits and o is in the upper
3343 bits. o_size is
3345 * 117: 4
3346 * 118: 6
3347 * 119: 5
3348 """
3349 gso_table = self._gso_table
3350 gso_df = self.df
3351 columns = list(gso_df.columns)
3352 selected = gso_df[self.columns]
3353 col_index = [(col, columns.index(col)) for col in self.columns]
3354 keys = np.empty(selected.shape, dtype=np.uint64)
3355 for o, (idx, row) in enumerate(selected.iterrows()):
3356 for j, (col, v) in enumerate(col_index):
3357 val = row[col]
3358 # Allow columns with mixed str and None or pd.NA (GH 23633)
3359 val = "" if isna(val) else val
3360 key = gso_table.get(val, None)
3361 if key is None:
3362 # Stata prefers human numbers
3363 key = (v + 1, o + 1)
3364 gso_table[val] = key
3365 keys[o, j] = self._convert_key(key)
3366 for i, col in enumerate(self.columns):
3367 gso_df[col] = keys[:, i]
3369 return gso_table, gso_df
3371 def generate_blob(self, gso_table: dict[str, tuple[int, int]]) -> bytes:
3372 """
3373 Generates the binary blob of GSOs that is written to the dta file.
3375 Parameters
3376 ----------
3377 gso_table : dict
3378 Ordered dictionary (str, vo)
3380 Returns
3381 -------
3382 gso : bytes
3383 Binary content of dta file to be placed between strl tags
3385 Notes
3386 -----
3387 Output format depends on dta version. 117 uses two uint32s to
3388 express v and o while 118+ uses a uint32 for v and a uint64 for o.
3389 """
3390 # Format information
3391 # Length includes null term
3392 # 117
3393 # GSOvvvvooootllllxxxxxxxxxxxxxxx...x
3394 # 3 u4 u4 u1 u4 string + null term
3395 #
3396 # 118, 119
3397 # GSOvvvvooooooootllllxxxxxxxxxxxxxxx...x
3398 # 3 u4 u8 u1 u4 string + null term
3400 bio = BytesIO()
3401 gso = bytes("GSO", "ascii")
3402 gso_type = struct.pack(self._byteorder + "B", 130)
3403 null = struct.pack(self._byteorder + "B", 0)
3404 v_type = self._byteorder + self._gso_v_type
3405 o_type = self._byteorder + self._gso_o_type
3406 len_type = self._byteorder + "I"
3407 for strl, vo in gso_table.items():
3408 if vo == (0, 0):
3409 continue
3410 v, o = vo
3412 # GSO
3413 bio.write(gso)
3415 # vvvv
3416 bio.write(struct.pack(v_type, v))
3418 # oooo / oooooooo
3419 bio.write(struct.pack(o_type, o))
3421 # t
3422 bio.write(gso_type)
3424 # llll
3425 if isinstance(strl, str):
3426 strl_convert = bytes(strl, "utf-8")
3427 else:
3428 strl_convert = strl
3430 bio.write(struct.pack(len_type, len(strl_convert) + 1))
3432 # xxx...xxx
3433 bio.write(strl_convert)
3434 bio.write(null)
3436 return bio.getvalue()
3439class StataWriter117(StataWriter):
3440 """
3441 A class for writing Stata binary dta files in Stata 13 format (117)
3443 Parameters
3444 ----------
3445 fname : path (string), buffer or path object
3446 string, pathlib.Path or
3447 object implementing a binary write() functions. If using a buffer
3448 then the buffer will not be automatically closed after the file
3449 is written.
3450 data : DataFrame
3451 Input to save
3452 convert_dates : dict
3453 Dictionary mapping columns containing datetime types to stata internal
3454 format to use when writing the dates. Options are 'tc', 'td', 'tm',
3455 'tw', 'th', 'tq', 'ty'. Column can be either an integer or a name.
3456 Datetime columns that do not have a conversion type specified will be
3457 converted to 'tc'. Raises NotImplementedError if a datetime column has
3458 timezone information
3459 write_index : bool
3460 Write the index to Stata dataset.
3461 byteorder : str
3462 Can be ">", "<", "little", or "big". default is `sys.byteorder`
3463 time_stamp : datetime
3464 A datetime to use as file creation date. Default is the current time
3465 data_label : str
3466 A label for the data set. Must be 80 characters or smaller.
3467 variable_labels : dict
3468 Dictionary containing columns as keys and variable labels as values.
3469 Each label must be 80 characters or smaller.
3470 convert_strl : list
3471 List of columns names to convert to Stata StrL format. Columns with
3472 more than 2045 characters are automatically written as StrL.
3473 Smaller columns can be converted by including the column name. Using
3474 StrLs can reduce output file size when strings are longer than 8
3475 characters, and either frequently repeated or sparse.
3476 {compression_options}
3478 value_labels : dict of dicts
3479 Dictionary containing columns as keys and dictionaries of column value
3480 to labels as values. The combined length of all labels for a single
3481 variable must be 32,000 characters or smaller.
3483 Returns
3484 -------
3485 writer : StataWriter117 instance
3486 The StataWriter117 instance has a write_file method, which will
3487 write the file to the given `fname`.
3489 Raises
3490 ------
3491 NotImplementedError
3492 * If datetimes contain timezone information
3493 ValueError
3494 * Columns listed in convert_dates are neither datetime64[ns]
3495 or datetime
3496 * Column dtype is not representable in Stata
3497 * Column listed in convert_dates is not in DataFrame
3498 * Categorical label contains more than 32,000 characters
3500 Examples
3501 --------
3502 >>> data = pd.DataFrame([[1.0, 1, "a"]], columns=["a", "b", "c"])
3503 >>> writer = pd.io.stata.StataWriter117("./data_file.dta", data)
3504 >>> writer.write_file()
3506 Directly write a zip file
3507 >>> compression = {"method": "zip", "archive_name": "data_file.dta"}
3508 >>> writer = pd.io.stata.StataWriter117(
3509 ... "./data_file.zip", data, compression=compression
3510 ... )
3511 >>> writer.write_file()
3513 Or with long strings stored in strl format
3514 >>> data = pd.DataFrame(
3515 ... [["A relatively long string"], [""], [""]], columns=["strls"]
3516 ... )
3517 >>> writer = pd.io.stata.StataWriter117(
3518 ... "./data_file_with_long_strings.dta", data, convert_strl=["strls"]
3519 ... )
3520 >>> writer.write_file()
3521 """
3523 _max_string_length = 2045
3524 _dta_version = 117
3526 def __init__(
3527 self,
3528 fname: FilePath | WriteBuffer[bytes],
3529 data: DataFrame,
3530 convert_dates: dict[Hashable, str] | None = None,
3531 write_index: bool = True,
3532 byteorder: str | None = None,
3533 time_stamp: datetime | None = None,
3534 data_label: str | None = None,
3535 variable_labels: dict[Hashable, str] | None = None,
3536 convert_strl: Sequence[Hashable] | None = None,
3537 compression: CompressionOptions = "infer",
3538 storage_options: StorageOptions | None = None,
3539 *,
3540 value_labels: dict[Hashable, dict[float, str]] | None = None,
3541 ) -> None:
3542 # Copy to new list since convert_strl might be modified later
3543 self._convert_strl: list[Hashable] = []
3544 if convert_strl is not None:
3545 self._convert_strl.extend(convert_strl)
3547 super().__init__(
3548 fname,
3549 data,
3550 convert_dates,
3551 write_index,
3552 byteorder=byteorder,
3553 time_stamp=time_stamp,
3554 data_label=data_label,
3555 variable_labels=variable_labels,
3556 value_labels=value_labels,
3557 compression=compression,
3558 storage_options=storage_options,
3559 )
3560 self._map: dict[str, int] = {}
3561 self._strl_blob = b""
3563 @staticmethod
3564 def _tag(val: str | bytes, tag: str) -> bytes:
3565 """Surround val with <tag></tag>"""
3566 if isinstance(val, str):
3567 val = bytes(val, "utf-8")
3568 return bytes("<" + tag + ">", "utf-8") + val + bytes("</" + tag + ">", "utf-8")
3570 def _update_map(self, tag: str) -> None:
3571 """Update map location for tag with file position"""
3572 assert self.handles.handle is not None
3573 self._map[tag] = self.handles.handle.tell()
3575 def _write_header(
3576 self,
3577 data_label: str | None = None,
3578 time_stamp: datetime | None = None,
3579 ) -> None:
3580 """Write the file header"""
3581 byteorder = self._byteorder
3582 self._write_bytes(bytes("<stata_dta>", "utf-8"))
3583 bio = BytesIO()
3584 # ds_format - 117
3585 bio.write(self._tag(bytes(str(self._dta_version), "utf-8"), "release"))
3586 # byteorder
3587 bio.write(self._tag((byteorder == ">" and "MSF") or "LSF", "byteorder"))
3588 # number of vars, 2 bytes in 117 and 118, 4 byte in 119
3589 nvar_type = "H" if self._dta_version <= 118 else "I"
3590 bio.write(self._tag(struct.pack(byteorder + nvar_type, self.nvar), "K"))
3591 # 117 uses 4 bytes, 118 uses 8
3592 nobs_size = "I" if self._dta_version == 117 else "Q"
3593 bio.write(self._tag(struct.pack(byteorder + nobs_size, self.nobs), "N"))
3594 # data label 81 bytes, char, null terminated
3595 label = data_label[:80] if data_label is not None else ""
3596 encoded_label = label.encode(self._encoding)
3597 label_size = "B" if self._dta_version == 117 else "H"
3598 label_len = struct.pack(byteorder + label_size, len(encoded_label))
3599 encoded_label = label_len + encoded_label
3600 bio.write(self._tag(encoded_label, "label"))
3601 # time stamp, 18 bytes, char, null terminated
3602 # format dd Mon yyyy hh:mm
3603 if time_stamp is None:
3604 time_stamp = datetime.now() # noqa: TID251
3605 elif not isinstance(time_stamp, datetime):
3606 raise ValueError("time_stamp should be datetime type")
3607 # Avoid locale-specific month conversion
3608 months = [
3609 "Jan",
3610 "Feb",
3611 "Mar",
3612 "Apr",
3613 "May",
3614 "Jun",
3615 "Jul",
3616 "Aug",
3617 "Sep",
3618 "Oct",
3619 "Nov",
3620 "Dec",
3621 ]
3622 month_lookup = {i + 1: month for i, month in enumerate(months)}
3623 ts = (
3624 time_stamp.strftime("%d ")
3625 + month_lookup[time_stamp.month]
3626 + time_stamp.strftime(" %Y %H:%M")
3627 )
3628 # '\x11' added due to inspection of Stata file
3629 stata_ts = b"\x11" + bytes(ts, "utf-8")
3630 bio.write(self._tag(stata_ts, "timestamp"))
3631 self._write_bytes(self._tag(bio.getvalue(), "header"))
3633 def _write_map(self) -> None:
3634 """
3635 Called twice during file write. The first populates the values in
3636 the map with 0s. The second call writes the final map locations when
3637 all blocks have been written.
3638 """
3639 if not self._map:
3640 self._map = {
3641 "stata_data": 0,
3642 "map": self.handles.handle.tell(),
3643 "variable_types": 0,
3644 "varnames": 0,
3645 "sortlist": 0,
3646 "formats": 0,
3647 "value_label_names": 0,
3648 "variable_labels": 0,
3649 "characteristics": 0,
3650 "data": 0,
3651 "strls": 0,
3652 "value_labels": 0,
3653 "stata_data_close": 0,
3654 "end-of-file": 0,
3655 }
3656 # Move to start of map
3657 self.handles.handle.seek(self._map["map"])
3658 bio = BytesIO()
3659 for val in self._map.values():
3660 bio.write(struct.pack(self._byteorder + "Q", val))
3661 self._write_bytes(self._tag(bio.getvalue(), "map"))
3663 def _write_variable_types(self) -> None:
3664 self._update_map("variable_types")
3665 bio = BytesIO()
3666 for typ in self.typlist:
3667 bio.write(struct.pack(self._byteorder + "H", typ))
3668 self._write_bytes(self._tag(bio.getvalue(), "variable_types"))
3670 def _write_varnames(self) -> None:
3671 self._update_map("varnames")
3672 bio = BytesIO()
3673 # 118 scales by 4 to accommodate utf-8 data worst case encoding
3674 vn_len = 32 if self._dta_version == 117 else 128
3675 for name in self.varlist:
3676 name = self._null_terminate_str(name)
3677 name = _pad_bytes_new(name[:32].encode(self._encoding), vn_len + 1)
3678 bio.write(name)
3679 self._write_bytes(self._tag(bio.getvalue(), "varnames"))
3681 def _write_sortlist(self) -> None:
3682 self._update_map("sortlist")
3683 sort_size = 2 if self._dta_version < 119 else 4
3684 self._write_bytes(self._tag(b"\x00" * sort_size * (self.nvar + 1), "sortlist"))
3686 def _write_formats(self) -> None:
3687 self._update_map("formats")
3688 bio = BytesIO()
3689 fmt_len = 49 if self._dta_version == 117 else 57
3690 for fmt in self.fmtlist:
3691 bio.write(_pad_bytes_new(fmt.encode(self._encoding), fmt_len))
3692 self._write_bytes(self._tag(bio.getvalue(), "formats"))
3694 def _write_value_label_names(self) -> None:
3695 self._update_map("value_label_names")
3696 bio = BytesIO()
3697 # 118 scales by 4 to accommodate utf-8 data worst case encoding
3698 vl_len = 32 if self._dta_version == 117 else 128
3699 for i in range(self.nvar):
3700 # Use variable name when categorical
3701 name = "" # default name
3702 if self._has_value_labels[i]:
3703 name = self.varlist[i]
3704 name = self._null_terminate_str(name)
3705 encoded_name = _pad_bytes_new(name[:32].encode(self._encoding), vl_len + 1)
3706 bio.write(encoded_name)
3707 self._write_bytes(self._tag(bio.getvalue(), "value_label_names"))
3709 def _write_variable_labels(self) -> None:
3710 # Missing labels are 80 blank characters plus null termination
3711 self._update_map("variable_labels")
3712 bio = BytesIO()
3713 # 118 scales by 4 to accommodate utf-8 data worst case encoding
3714 vl_len = 80 if self._dta_version == 117 else 320
3715 blank = _pad_bytes_new("", vl_len + 1)
3717 if self._variable_labels is None:
3718 for _ in range(self.nvar):
3719 bio.write(blank)
3720 self._write_bytes(self._tag(bio.getvalue(), "variable_labels"))
3721 return
3723 for col in self.data:
3724 if col in self._variable_labels:
3725 label = self._variable_labels[col]
3726 if len(label) > 80:
3727 raise ValueError("Variable labels must be 80 characters or fewer")
3728 try:
3729 encoded = label.encode(self._encoding)
3730 except UnicodeEncodeError as err:
3731 raise ValueError(
3732 "Variable labels must contain only characters that "
3733 f"can be encoded in {self._encoding}"
3734 ) from err
3736 bio.write(_pad_bytes_new(encoded, vl_len + 1))
3737 else:
3738 bio.write(blank)
3739 self._write_bytes(self._tag(bio.getvalue(), "variable_labels"))
3741 def _write_characteristics(self) -> None:
3742 self._update_map("characteristics")
3743 self._write_bytes(self._tag(b"", "characteristics"))
3745 def _write_data(self, records: np.rec.recarray) -> None:
3746 self._update_map("data")
3747 self._write_bytes(b"<data>")
3748 self._write_bytes(records.tobytes())
3749 self._write_bytes(b"</data>")
3751 def _write_strls(self) -> None:
3752 self._update_map("strls")
3753 self._write_bytes(self._tag(self._strl_blob, "strls"))
3755 def _write_expansion_fields(self) -> None:
3756 """No-op in dta 117+"""
3758 def _write_value_labels(self) -> None:
3759 self._update_map("value_labels")
3760 bio = BytesIO()
3761 for vl in self._value_labels:
3762 lab = vl.generate_value_label(self._byteorder)
3763 lab = self._tag(lab, "lbl")
3764 bio.write(lab)
3765 self._write_bytes(self._tag(bio.getvalue(), "value_labels"))
3767 def _write_file_close_tag(self) -> None:
3768 self._update_map("stata_data_close")
3769 self._write_bytes(bytes("</stata_dta>", "utf-8"))
3770 self._update_map("end-of-file")
3772 def _update_strl_names(self) -> None:
3773 """
3774 Update column names for conversion to strl if they might have been
3775 changed to comply with Stata naming rules
3776 """
3777 # Update convert_strl if names changed
3778 for orig, new in self._converted_names.items():
3779 if orig in self._convert_strl:
3780 idx = self._convert_strl.index(orig)
3781 self._convert_strl[idx] = new
3783 def _convert_strls(self, data: DataFrame) -> DataFrame:
3784 """
3785 Convert columns to StrLs if either very large or in the
3786 convert_strl variable
3787 """
3788 convert_cols = [
3789 col
3790 for i, col in enumerate(data)
3791 if self.typlist[i] == 32768 or col in self._convert_strl
3792 ]
3794 if convert_cols:
3795 ssw = StataStrLWriter(
3796 data, convert_cols, version=self._dta_version, byteorder=self._byteorder
3797 )
3798 tab, new_data = ssw.generate_table()
3799 data = new_data
3800 self._strl_blob = ssw.generate_blob(tab)
3801 return data
3803 def _set_formats_and_types(self, dtypes: Series) -> None:
3804 self.typlist = []
3805 self.fmtlist = []
3806 for col, dtype in dtypes.items():
3807 force_strl = col in self._convert_strl
3808 fmt = _dtype_to_default_stata_fmt(
3809 dtype,
3810 self.data[col],
3811 dta_version=self._dta_version,
3812 force_strl=force_strl,
3813 )
3814 self.fmtlist.append(fmt)
3815 self.typlist.append(
3816 _dtype_to_stata_type_117(dtype, self.data[col], force_strl)
3817 )
3820class StataWriterUTF8(StataWriter117):
3821 """
3822 Stata binary dta file writing in Stata 15 (118) and 16 (119) formats
3824 DTA 118 and 119 format files support unicode string data (both fixed
3825 and strL) format. Unicode is also supported in value labels, variable
3826 labels and the dataset label. Format 119 is automatically used if the
3827 file contains more than 32,767 variables.
3829 Parameters
3830 ----------
3831 fname : path (string), buffer or path object
3832 string, pathlib.Path or
3833 object implementing a binary write() functions. If using a buffer
3834 then the buffer will not be automatically closed after the file
3835 is written.
3836 data : DataFrame
3837 Input to save
3838 convert_dates : dict, default None
3839 Dictionary mapping columns containing datetime types to stata internal
3840 format to use when writing the dates. Options are 'tc', 'td', 'tm',
3841 'tw', 'th', 'tq', 'ty'. Column can be either an integer or a name.
3842 Datetime columns that do not have a conversion type specified will be
3843 converted to 'tc'. Raises NotImplementedError if a datetime column has
3844 timezone information
3845 write_index : bool, default True
3846 Write the index to Stata dataset.
3847 byteorder : str, default None
3848 Can be ">", "<", "little", or "big". default is `sys.byteorder`
3849 time_stamp : datetime, default None
3850 A datetime to use as file creation date. Default is the current time
3851 data_label : str, default None
3852 A label for the data set. Must be 80 characters or smaller.
3853 variable_labels : dict, default None
3854 Dictionary containing columns as keys and variable labels as values.
3855 Each label must be 80 characters or smaller.
3856 convert_strl : list, default None
3857 List of columns names to convert to Stata StrL format. Columns with
3858 more than 2045 characters are automatically written as StrL.
3859 Smaller columns can be converted by including the column name. Using
3860 StrLs can reduce output file size when strings are longer than 8
3861 characters, and either frequently repeated or sparse.
3862 version : int, default None
3863 The dta version to use. By default, uses the size of data to determine
3864 the version. 118 is used if data.shape[1] <= 32767, and 119 is used
3865 for storing larger DataFrames.
3866 {compression_options}
3868 value_labels : dict of dicts
3869 Dictionary containing columns as keys and dictionaries of column value
3870 to labels as values. The combined length of all labels for a single
3871 variable must be 32,000 characters or smaller.
3873 Returns
3874 -------
3875 StataWriterUTF8
3876 The instance has a write_file method, which will write the file to the
3877 given `fname`.
3879 Raises
3880 ------
3881 NotImplementedError
3882 * If datetimes contain timezone information
3883 ValueError
3884 * Columns listed in convert_dates are neither datetime64[ns]
3885 or datetime
3886 * Column dtype is not representable in Stata
3887 * Column listed in convert_dates is not in DataFrame
3888 * Categorical label contains more than 32,000 characters
3890 Examples
3891 --------
3892 Using Unicode data and column names
3894 >>> from pandas.io.stata import StataWriterUTF8
3895 >>> data = pd.DataFrame([[1.0, 1, "ᴬ"]], columns=["a", "β", "ĉ"])
3896 >>> writer = StataWriterUTF8("./data_file.dta", data)
3897 >>> writer.write_file()
3899 Directly write a zip file
3900 >>> compression = {"method": "zip", "archive_name": "data_file.dta"}
3901 >>> writer = StataWriterUTF8("./data_file.zip", data, compression=compression)
3902 >>> writer.write_file()
3904 Or with long strings stored in strl format
3906 >>> data = pd.DataFrame(
3907 ... [["ᴀ relatively long ŝtring"], [""], [""]], columns=["strls"]
3908 ... )
3909 >>> writer = StataWriterUTF8(
3910 ... "./data_file_with_long_strings.dta", data, convert_strl=["strls"]
3911 ... )
3912 >>> writer.write_file()
3913 """
3915 _encoding: Literal["utf-8"] = "utf-8"
3917 def __init__(
3918 self,
3919 fname: FilePath | WriteBuffer[bytes],
3920 data: DataFrame,
3921 convert_dates: dict[Hashable, str] | None = None,
3922 write_index: bool = True,
3923 byteorder: str | None = None,
3924 time_stamp: datetime | None = None,
3925 data_label: str | None = None,
3926 variable_labels: dict[Hashable, str] | None = None,
3927 convert_strl: Sequence[Hashable] | None = None,
3928 version: int | None = None,
3929 compression: CompressionOptions = "infer",
3930 storage_options: StorageOptions | None = None,
3931 *,
3932 value_labels: dict[Hashable, dict[float, str]] | None = None,
3933 ) -> None:
3934 if version is None:
3935 version = 118 if data.shape[1] <= 32767 else 119
3936 elif version not in (118, 119):
3937 raise ValueError("version must be either 118 or 119.")
3938 elif version == 118 and data.shape[1] > 32767:
3939 raise ValueError(
3940 "You must use version 119 for data sets containing more than"
3941 "32,767 variables"
3942 )
3944 super().__init__(
3945 fname,
3946 data,
3947 convert_dates=convert_dates,
3948 write_index=write_index,
3949 byteorder=byteorder,
3950 time_stamp=time_stamp,
3951 data_label=data_label,
3952 variable_labels=variable_labels,
3953 value_labels=value_labels,
3954 convert_strl=convert_strl,
3955 compression=compression,
3956 storage_options=storage_options,
3957 )
3958 # Override version set in StataWriter117 init
3959 self._dta_version = version
3961 def _validate_variable_name(self, name: str) -> str:
3962 """
3963 Validate variable names for Stata export.
3965 Parameters
3966 ----------
3967 name : str
3968 Variable name
3970 Returns
3971 -------
3972 str
3973 The validated name with invalid characters replaced with
3974 underscores.
3976 Notes
3977 -----
3978 Stata 118+ support most unicode characters. The only limitation is in
3979 the ascii range where the characters supported are a-z, A-Z, 0-9 and _.
3980 """
3981 # High code points appear to be acceptable
3982 for c in name:
3983 if (
3984 (
3985 ord(c) < 128
3986 and (c < "A" or c > "Z")
3987 and (c < "a" or c > "z")
3988 and (c < "0" or c > "9")
3989 and c != "_"
3990 )
3991 or 128 <= ord(c) < 192
3992 or c in {"×", "÷"} # noqa: RUF001
3993 ):
3994 name = name.replace(c, "_")
3996 return name