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