TYP: pandas/io/parquet.py (#29972)

Author: simonjayhawkins

pandas/io/parquet.py

@@ -1,5 +1,6 @@
 """ parquet compat """
 
+from typing import Any, Dict, Optional
 from warnings import catch_warnings
 
 from pandas.compat._optional import import_optional_dependency
@@ -10,7 +11,7 @@
 from pandas.io.common import get_filepath_or_buffer, is_gcs_url, is_s3_url
 
 
-def get_engine(engine):
+def get_engine(engine: str) -> "BaseImpl":
     """ return our implementation """
 
     if engine == "auto":
@@ -35,19 +36,15 @@ def get_engine(engine):
             "support"
         )
 
-    if engine not in ["pyarrow", "fastparquet"]:
-        raise ValueError("engine must be one of 'pyarrow', 'fastparquet'")
-
     if engine == "pyarrow":
         return PyArrowImpl()
     elif engine == "fastparquet":
         return FastParquetImpl()
 
+    raise ValueError("engine must be one of 'pyarrow', 'fastparquet'")
 
-class BaseImpl:
-
-    api = None  # module
 
+class BaseImpl:
     @staticmethod
     def validate_dataframe(df):
 
@@ -74,7 +71,7 @@ def read(self, path, columns=None, **kwargs):
 
 class PyArrowImpl(BaseImpl):
     def __init__(self):
-        pyarrow = import_optional_dependency(
+        import_optional_dependency(
             "pyarrow", extra="pyarrow is required for parquet support."
         )
         import pyarrow.parquet
@@ -87,13 +84,14 @@ def write(
         path,
         compression="snappy",
         coerce_timestamps="ms",
-        index=None,
+        index: Optional[bool] = None,
         partition_cols=None,
         **kwargs,
     ):
         self.validate_dataframe(df)
         path, _, _, _ = get_filepath_or_buffer(path, mode="wb")
 
+        from_pandas_kwargs: Dict[str, Any]
         if index is None:
             from_pandas_kwargs = {}
         else:
@@ -203,7 +201,7 @@ def to_parquet(
     path,
     engine="auto",
     compression="snappy",
-    index=None,
+    index: Optional[bool] = None,
     partition_cols=None,
     **kwargs,
 ):