CLN: Restructure to_stata in frame

Restructure to get mypy to like it
pandas-dev · Jan 16, 2020 · 8043174 · 8043174
1 parent 335a3e4
commit 8043174
Showing 1 changed file with 44 additions and 27 deletions.
diff --git a/pandas/core/frame.py b/pandas/core/frame.py
@@ -1951,35 +1951,52 @@ 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:
-            if convert_strl is not None:
-                raise ValueError("strl is not supported in format 114")
-            from pandas.io.stata import StataWriter as statawriter
+            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()
         elif version == 117:
-            from pandas.io.stata import StataWriter117 as statawriter
-        else:  # versions 118 and 119
-            from pandas.io.stata import StataWriterUTF8 as statawriter
-
-        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)
-            kwargs["version"] = version
-
-        writer = 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,
-            **kwargs,
-        )
-        writer.write_file()
+            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()
 
     @deprecate_kwarg(old_arg_name="fname", new_arg_name="path")
     def to_feather(self, path) -> None: