TYP: stricter type for compression (#44854)

Author: twoertwein

pandas/_typing.py

@@ -242,7 +242,9 @@ def closed(self) -> bool:
 
 # compression keywords and compression
 CompressionDict = Dict[str, Any]
-CompressionOptions = Optional[Union[str, CompressionDict]]
+CompressionOptions = Optional[
+    Union[Literal["infer", "gzip", "bz2", "zip", "xz"], CompressionDict]
+]
 
 
 # types in DataFrameFormatter

pandas/io/parsers/readers.py

@@ -19,6 +19,7 @@
 from pandas._libs.parsers import STR_NA_VALUES
 from pandas._typing import (
     ArrayLike,
+    CompressionOptions,
     DtypeArg,
     FilePath,
     ReadCsvBuffer,
@@ -618,7 +619,7 @@ def read_csv(
     iterator=False,
     chunksize=None,
     # Quoting, Compression, and File Format
-    compression="infer",
+    compression: CompressionOptions = "infer",
     thousands=None,
     decimal: str = ".",
     lineterminator=None,
@@ -716,7 +717,7 @@ def read_table(
     iterator=False,
     chunksize=None,
     # Quoting, Compression, and File Format
-    compression="infer",
+    compression: CompressionOptions = "infer",
     thousands=None,
     decimal: str = ".",
     lineterminator=None,

pandas/io/xml.py

@@ -105,8 +105,8 @@ def __init__(
         names,
         encoding,
         stylesheet,
-        compression,
-        storage_options,
+        compression: CompressionOptions,
+        storage_options: StorageOptions,
     ) -> None:
         self.path_or_buffer = path_or_buffer
         self.xpath = xpath
@@ -570,8 +570,8 @@ def _transform_doc(self) -> bytes:
 def get_data_from_filepath(
     filepath_or_buffer: FilePath | bytes | ReadBuffer[bytes] | ReadBuffer[str],
     encoding,
-    compression,
-    storage_options,
+    compression: CompressionOptions,
+    storage_options: StorageOptions,
 ) -> str | bytes | ReadBuffer[bytes] | ReadBuffer[str]:
     """
     Extract raw XML data.
@@ -666,8 +666,8 @@ def _parse(
     encoding,
     parser,
     stylesheet,
-    compression,
-    storage_options,
+    compression: CompressionOptions,
+    storage_options: StorageOptions,
     **kwargs,
 ) -> DataFrame:
     """

pandas/tests/io/xml/test_to_xml.py

@@ -1311,7 +1311,12 @@ def test_filename_and_suffix_comp(parser, comp, compfile):
 def test_unsuported_compression(datapath, parser):
     with pytest.raises(ValueError, match="Unrecognized compression type"):
         with tm.ensure_clean() as path:
-            geom_df.to_xml(path, parser=parser, compression="7z")
+            # Argument "compression" to "to_xml" of "DataFrame" has incompatible type
+            # "Literal['7z']"; expected "Union[Literal['infer'], Literal['gzip'],
+            # Literal['bz2'], Literal['zip'], Literal['xz'], Dict[str, Any], None]"
+            geom_df.to_xml(
+                path, parser=parser, compression="7z"  # type: ignore[arg-type]
+            )
 
 
 # STORAGE OPTIONS

pandas/tests/io/xml/test_xml.py

@@ -1069,7 +1069,10 @@ def test_wrong_compression_zip(parser, comp):
 def test_unsuported_compression(datapath, parser):
     with pytest.raises(ValueError, match="Unrecognized compression type"):
         with tm.ensure_clean() as path:
-            read_xml(path, parser=parser, compression="7z")
+            # error: Argument "compression" to "read_xml" has incompatible type
+            # "Literal['7z']"; expected "Union[Literal['infer'], Literal['gzip'],
+            # Literal['bz2'], Literal['zip'], Literal['xz'], Dict[str, Any], None]"
+            read_xml(path, parser=parser, compression="7z")  # type: ignore[arg-type]
 
 
 # STORAGE OPTIONS