Skip to content

Commit e341df2

Browse files
simonjayhawkinsjreback
authored andcommitted
TYP: pandas/io/parquet.py (#29972)
1 parent 83812e1 commit e341df2

File tree

1 file changed

+8
-10
lines changed

1 file changed

+8
-10
lines changed

pandas/io/parquet.py

Lines changed: 8 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,6 @@
11
""" parquet compat """
22

3+
from typing import Any, Dict, Optional
34
from warnings import catch_warnings
45

56
from pandas.compat._optional import import_optional_dependency
@@ -10,7 +11,7 @@
1011
from pandas.io.common import get_filepath_or_buffer, is_gcs_url, is_s3_url
1112

1213

13-
def get_engine(engine):
14+
def get_engine(engine: str) -> "BaseImpl":
1415
""" return our implementation """
1516

1617
if engine == "auto":
@@ -35,19 +36,15 @@ def get_engine(engine):
3536
"support"
3637
)
3738

38-
if engine not in ["pyarrow", "fastparquet"]:
39-
raise ValueError("engine must be one of 'pyarrow', 'fastparquet'")
40-
4139
if engine == "pyarrow":
4240
return PyArrowImpl()
4341
elif engine == "fastparquet":
4442
return FastParquetImpl()
4543

44+
raise ValueError("engine must be one of 'pyarrow', 'fastparquet'")
4645

47-
class BaseImpl:
48-
49-
api = None # module
5046

47+
class BaseImpl:
5148
@staticmethod
5249
def validate_dataframe(df):
5350

@@ -74,7 +71,7 @@ def read(self, path, columns=None, **kwargs):
7471

7572
class PyArrowImpl(BaseImpl):
7673
def __init__(self):
77-
pyarrow = import_optional_dependency(
74+
import_optional_dependency(
7875
"pyarrow", extra="pyarrow is required for parquet support."
7976
)
8077
import pyarrow.parquet
@@ -87,13 +84,14 @@ def write(
8784
path,
8885
compression="snappy",
8986
coerce_timestamps="ms",
90-
index=None,
87+
index: Optional[bool] = None,
9188
partition_cols=None,
9289
**kwargs,
9390
):
9491
self.validate_dataframe(df)
9592
path, _, _, _ = get_filepath_or_buffer(path, mode="wb")
9693

94+
from_pandas_kwargs: Dict[str, Any]
9795
if index is None:
9896
from_pandas_kwargs = {}
9997
else:
@@ -203,7 +201,7 @@ def to_parquet(
203201
path,
204202
engine="auto",
205203
compression="snappy",
206-
index=None,
204+
index: Optional[bool] = None,
207205
partition_cols=None,
208206
**kwargs,
209207
):

0 commit comments

Comments
 (0)