ENH: Add literals

Add literals for limited value inputs
pandas-dev · Aug 18, 2022 · 620f297 · 620f297
1 parent 71fc68a
commit 620f297
Show file tree

Hide file tree

Showing 3 changed files with 32 additions and 13 deletions.
diff --git a/pandas-stubs/_typing.pyi b/pandas-stubs/_typing.pyi
@@ -203,4 +203,21 @@ GroupByObjectNonScalar = Union[
 ]
 GroupByObject = Union[Scalar, GroupByObjectNonScalar]
 
+StataDateFormat = Literal[
+    "tc",
+    "%tc",
+    "td",
+    "%td",
+    "tw",
+    "%tw",
+    "tm",
+    "%tm",
+    "tq",
+    "%tq",
+    "th",
+    "%th",
+    "ty",
+    "%ty",
+]
+
 __all__ = ["npt", "type_t"]
diff --git a/pandas-stubs/core/frame.pyi b/pandas-stubs/core/frame.pyi
@@ -69,6 +69,7 @@ from pandas._typing import (
     Scalar,
     ScalarT,
     SeriesAxisType,
+    StataDateFormat,
     StorageOptions,
     StrLike,
     T as TType,
@@ -241,13 +242,13 @@ class DataFrame(NDFrame, OpsMixin):
     def to_stata(
         self,
         path: FilePath | WriteBuffer[bytes],
-        convert_dates: dict[Hashable, str] | None = ...,
+        convert_dates: dict[HashableT, StataDateFormat] | None = ...,
         write_index: _bool = ...,
         byteorder: Literal["<", ">", "little", "big"] | None = ...,
         time_stamp: _dt.datetime | None = ...,
         data_label: _str | None = ...,
-        variable_labels: dict[Hashable, str] | None = ...,
-        version: int | None = ...,
+        variable_labels: dict[HashableT, str] | None = ...,
+        version: Literal[114, 117, 118, 119] | None = ...,
         convert_strl: list[HashableT] | None = ...,
         compression: CompressionOptions = ...,
         storage_options: StorageOptions = ...,

diff --git a/pandas-stubs/io/stata.pyi b/pandas-stubs/io/stata.pyi
@@ -19,6 +19,7 @@ from pandas._typing import (
     FilePath,
     HashableT,
     ReadBuffer,
+    StataDateFormat,
     StorageOptions,
     WriteBuffer,
 )
@@ -138,16 +139,16 @@ class StataWriter(StataParser):
         self,
         fname: FilePath | WriteBuffer[bytes],
         data: DataFrame,
-        convert_dates: dict[Hashable, str] | None = ...,
+        convert_dates: dict[HashableT, StataDateFormat] | None = ...,
         write_index: bool = ...,
         byteorder: str | None = ...,
         time_stamp: datetime.datetime | None = ...,
         data_label: str | None = ...,
-        variable_labels: dict[Hashable, str] | None = ...,
+        variable_labels: dict[HashableT, str] | None = ...,
         compression: CompressionOptions = ...,
         storage_options: StorageOptions = ...,
         *,
-        value_labels: dict[Hashable, dict[float, str]] | None = ...,
+        value_labels: dict[HashableT, dict[float, str]] | None = ...,
     ) -> None: ...
     def write_file(self) -> None: ...
 
@@ -156,34 +157,34 @@ class StataWriter117(StataWriter):
         self,
         fname: FilePath | WriteBuffer[bytes],
         data: DataFrame,
-        convert_dates: dict[Hashable, str] | None = ...,
+        convert_dates: dict[HashableT, StataDateFormat] | None = ...,
         write_index: bool = ...,
         byteorder: str | None = ...,
         time_stamp: datetime.datetime | None = ...,
         data_label: str | None = ...,
-        variable_labels: dict[Hashable, str] | None = ...,
+        variable_labels: dict[HashableT, str] | None = ...,
         convert_strl: Sequence[Hashable] | None = ...,
         compression: CompressionOptions = ...,
         storage_options: StorageOptions = ...,
         *,
-        value_labels: dict[Hashable, dict[float, str]] | None = ...,
+        value_labels: dict[HashableT, dict[float, str]] | None = ...,
     ) -> None: ...
 
 class StataWriterUTF8(StataWriter117):
     def __init__(
         self,
         fname: FilePath | WriteBuffer[bytes],
         data: DataFrame,
-        convert_dates: dict[Hashable, str] | None = ...,
+        convert_dates: dict[HashableT, StataDateFormat] | None = ...,
         write_index: bool = ...,
         byteorder: str | None = ...,
         time_stamp: datetime.datetime | None = ...,
         data_label: str | None = ...,
-        variable_labels: dict[Hashable, str] | None = ...,
+        variable_labels: dict[HashableT, str] | None = ...,
         convert_strl: Sequence[Hashable] | None = ...,
-        version: int | None = ...,
+        version: Literal[118, 119] | None = ...,
         compression: CompressionOptions = ...,
         storage_options: StorageOptions = ...,
         *,
-        value_labels: dict[Hashable, dict[float, str]] | None = ...,
+        value_labels: dict[HashableT, dict[float, str]] | None = ...,
     ) -> None: ...