CLN: Revert frame

Revert frame
Add some annotations to help mypy

Author: bashtage

pandas/core/frame.py

@@ -1948,52 +1948,39 @@ def to_stata(
         """
         if version not in (114, 117, 118, 119, None):
             raise ValueError("Only formats 114, 117, 118 and 119 are supported.")
-        if version == 114 and convert_strl is not None:
-            raise ValueError("strl is not supported in format 114")
-
-        # TODO: There must be a better way?
         if version == 114:
-            from pandas.io.stata import StataWriter
-
-            StataWriter(
-                path,
-                self,
-                convert_dates=convert_dates,
-                byteorder=byteorder,
-                time_stamp=time_stamp,
-                data_label=data_label,
-                write_index=write_index,
-                variable_labels=variable_labels,
-            ).write_file()
+            if convert_strl is not None:
+                raise ValueError("strl is not supported in format 114")
+            from pandas.io.stata import StataWriter as statawriter
         elif version == 117:
-            from pandas.io.stata import StataWriter117
-
-            StataWriter117(
-                path,
-                self,
-                convert_dates=convert_dates,
-                byteorder=byteorder,
-                time_stamp=time_stamp,
-                data_label=data_label,
-                write_index=write_index,
-                variable_labels=variable_labels,
-                convert_strl=convert_strl,
-            ).write_file()
-        else:
-            from pandas.io.stata import StataWriterUTF8
-
-            StataWriterUTF8(
-                path,
-                self,
-                convert_dates=convert_dates,
-                byteorder=byteorder,
-                time_stamp=time_stamp,
-                data_label=data_label,
-                write_index=write_index,
-                variable_labels=variable_labels,
-                convert_strl=convert_strl,
-                version=version,
-            ).write_file()
+            # mypy: Name 'statawriter' already defined (possibly by an import)
+            from pandas.io.stata import StataWriter117 as statawriter  # type: ignore
+        else:  # versions 118 and 119
+            # mypy: Name 'statawriter' already defined (possibly by an import)
+            from pandas.io.stata import StataWriterUTF8 as statawriter  # type:ignore
+
+        kwargs = {}
+        if version is None or version >= 117:
+            # strl conversion is only supported >= 117
+            kwargs["convert_strl"] = convert_strl
+        if version is None or version >= 118:
+            # Specifying the version is only supported for UTF8 (118 or 119)
+            # mypy: Incompatible types in assignment
+            kwargs["version"] = version  # type: ignore
+
+        # mypy: Too many arguments for "StataWriter"
+        writer = statawriter(  # type: ignore
+            path,
+            self,
+            convert_dates=convert_dates,
+            byteorder=byteorder,
+            time_stamp=time_stamp,
+            data_label=data_label,
+            write_index=write_index,
+            variable_labels=variable_labels,
+            **kwargs,
+        )
+        writer.write_file()
 
     @deprecate_kwarg(old_arg_name="fname", new_arg_name="path")
     def to_feather(self, path) -> None:

pandas/io/stata.py

@@ -309,21 +309,21 @@ def convert_delta_safe(
     # Delta days relative to base
     elif fmt.startswith(("%td", "td", "%d", "d")):
         base = stata_epoch
-        days = dates
+        days: Series = dates
         conv_dates = convert_delta_safe(base, days, "d")
     # does not count leap days - 7 days is a week.
     # 52nd week may have more than 7 days
     elif fmt.startswith(("%tw", "tw")):
-        year = stata_epoch.year + dates // 52
+        year: Series = stata_epoch.year + dates // 52
         days = (dates % 52) * 7
         conv_dates = convert_year_days_safe(year, days)
     elif fmt.startswith(("%tm", "tm")):  # Delta months relative to base
         year = stata_epoch.year + dates // 12
-        month = (dates % 12) + 1
+        month: Series = (dates % 12) + 1
         conv_dates = convert_year_month_safe(year, month)
     elif fmt.startswith(("%tq", "tq")):  # Delta quarters relative to base
         year = stata_epoch.year + dates // 4
-        quarter_month = (dates % 4) * 3 + 1  # type: Series
+        quarter_month: Series = (dates % 4) * 3 + 1
         conv_dates = convert_year_month_safe(year, quarter_month)
     elif fmt.startswith(("%th", "th")):  # Delta half-years relative to base
         year = stata_epoch.year + dates // 2
@@ -361,19 +361,19 @@ def _datetime_to_stata_elapsed_vec(dates: Series, fmt: str) -> Series:
     def parse_dates_safe(
         dates: Series, delta: bool = False, year: bool = False, days: bool = False
     ):
-        d = {}
+        d: Dict[str, Any] = {}
         if is_datetime64_dtype(dates.values):
             if delta:
-                delta = dates - stata_epoch
-                d["delta"] = delta.values.astype(np.int64) // 1000  # microseconds
+                time_delta: Series = dates - stata_epoch
+                d["delta"] = time_delta.values.astype(np.int64) // 1000  # microseconds
             if days or year:
                 date_index = DatetimeIndex(dates)
                 d["year"], d["month"] = date_index.year, date_index.month
             if days:
-                days = dates.astype(np.int64) - to_datetime(
+                days_in_ns: Series = dates.astype(np.int64) - to_datetime(
                     d["year"], format="%Y"
                 ).astype(np.int64)
-                d["days"] = days // NS_PER_DAY
+                d["days"] = days_in_ns // NS_PER_DAY
 
         elif infer_dtype(dates, skipna=False) == "datetime":
             if delta: