Skip to content

Commit f10a282

Browse files
committed
CLN: Revert frame
Revert frame Add some annotations to help mypy
1 parent 6fc3a4b commit f10a282

File tree

2 files changed

+40
-53
lines changed

2 files changed

+40
-53
lines changed

pandas/core/frame.py

Lines changed: 31 additions & 44 deletions
Original file line numberDiff line numberDiff line change
@@ -1951,52 +1951,39 @@ def to_stata(
19511951
"""
19521952
if version not in (114, 117, 118, 119, None):
19531953
raise ValueError("Only formats 114, 117, 118 and 119 are supported.")
1954-
if version == 114 and convert_strl is not None:
1955-
raise ValueError("strl is not supported in format 114")
1956-
1957-
# TODO: There must be a better way?
19581954
if version == 114:
1959-
from pandas.io.stata import StataWriter
1960-
1961-
StataWriter(
1962-
path,
1963-
self,
1964-
convert_dates=convert_dates,
1965-
byteorder=byteorder,
1966-
time_stamp=time_stamp,
1967-
data_label=data_label,
1968-
write_index=write_index,
1969-
variable_labels=variable_labels,
1970-
).write_file()
1955+
if convert_strl is not None:
1956+
raise ValueError("strl is not supported in format 114")
1957+
from pandas.io.stata import StataWriter as statawriter
19711958
elif version == 117:
1972-
from pandas.io.stata import StataWriter117
1973-
1974-
StataWriter117(
1975-
path,
1976-
self,
1977-
convert_dates=convert_dates,
1978-
byteorder=byteorder,
1979-
time_stamp=time_stamp,
1980-
data_label=data_label,
1981-
write_index=write_index,
1982-
variable_labels=variable_labels,
1983-
convert_strl=convert_strl,
1984-
).write_file()
1985-
else:
1986-
from pandas.io.stata import StataWriterUTF8
1987-
1988-
StataWriterUTF8(
1989-
path,
1990-
self,
1991-
convert_dates=convert_dates,
1992-
byteorder=byteorder,
1993-
time_stamp=time_stamp,
1994-
data_label=data_label,
1995-
write_index=write_index,
1996-
variable_labels=variable_labels,
1997-
convert_strl=convert_strl,
1998-
version=version,
1999-
).write_file()
1959+
# mypy: Name 'statawriter' already defined (possibly by an import)
1960+
from pandas.io.stata import StataWriter117 as statawriter # type: ignore
1961+
else: # versions 118 and 119
1962+
# mypy: Name 'statawriter' already defined (possibly by an import)
1963+
from pandas.io.stata import StataWriterUTF8 as statawriter # type:ignore
1964+
1965+
kwargs = {}
1966+
if version is None or version >= 117:
1967+
# strl conversion is only supported >= 117
1968+
kwargs["convert_strl"] = convert_strl
1969+
if version is None or version >= 118:
1970+
# Specifying the version is only supported for UTF8 (118 or 119)
1971+
# mypy: Incompatible types in assignment
1972+
kwargs["version"] = version # type: ignore
1973+
1974+
# mypy: Too many arguments for "StataWriter"
1975+
writer = statawriter( # type: ignore
1976+
path,
1977+
self,
1978+
convert_dates=convert_dates,
1979+
byteorder=byteorder,
1980+
time_stamp=time_stamp,
1981+
data_label=data_label,
1982+
write_index=write_index,
1983+
variable_labels=variable_labels,
1984+
**kwargs,
1985+
)
1986+
writer.write_file()
20001987

20011988
@deprecate_kwarg(old_arg_name="fname", new_arg_name="path")
20021989
def to_feather(self, path) -> None:

pandas/io/stata.py

Lines changed: 9 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -309,21 +309,21 @@ def convert_delta_safe(
309309
# Delta days relative to base
310310
elif fmt.startswith(("%td", "td", "%d", "d")):
311311
base = stata_epoch
312-
days = dates
312+
days: Series = dates
313313
conv_dates = convert_delta_safe(base, days, "d")
314314
# does not count leap days - 7 days is a week.
315315
# 52nd week may have more than 7 days
316316
elif fmt.startswith(("%tw", "tw")):
317-
year = stata_epoch.year + dates // 52
317+
year: Series = stata_epoch.year + dates // 52
318318
days = (dates % 52) * 7
319319
conv_dates = convert_year_days_safe(year, days)
320320
elif fmt.startswith(("%tm", "tm")): # Delta months relative to base
321321
year = stata_epoch.year + dates // 12
322-
month = (dates % 12) + 1
322+
month: Series = (dates % 12) + 1
323323
conv_dates = convert_year_month_safe(year, month)
324324
elif fmt.startswith(("%tq", "tq")): # Delta quarters relative to base
325325
year = stata_epoch.year + dates // 4
326-
quarter_month = (dates % 4) * 3 + 1 # type: Series
326+
quarter_month: Series = (dates % 4) * 3 + 1
327327
conv_dates = convert_year_month_safe(year, quarter_month)
328328
elif fmt.startswith(("%th", "th")): # Delta half-years relative to base
329329
year = stata_epoch.year + dates // 2
@@ -361,19 +361,19 @@ def _datetime_to_stata_elapsed_vec(dates: Series, fmt: str) -> Series:
361361
def parse_dates_safe(
362362
dates: Series, delta: bool = False, year: bool = False, days: bool = False
363363
):
364-
d = {}
364+
d: Dict[str, Any] = {}
365365
if is_datetime64_dtype(dates.values):
366366
if delta:
367-
delta = dates - stata_epoch
368-
d["delta"] = delta.values.astype(np.int64) // 1000 # microseconds
367+
time_delta: Series = dates - stata_epoch
368+
d["delta"] = time_delta.values.astype(np.int64) // 1000 # microseconds
369369
if days or year:
370370
date_index = DatetimeIndex(dates)
371371
d["year"], d["month"] = date_index.year, date_index.month
372372
if days:
373-
days = dates.astype(np.int64) - to_datetime(
373+
days_in_ns: Series = dates.astype(np.int64) - to_datetime(
374374
d["year"], format="%Y"
375375
).astype(np.int64)
376-
d["days"] = days // NS_PER_DAY
376+
d["days"] = days_in_ns // NS_PER_DAY
377377

378378
elif infer_dtype(dates, skipna=False) == "datetime":
379379
if delta:

0 commit comments

Comments
 (0)