diff --git a/pandas/io/sas/sasreader.py b/pandas/io/sas/sasreader.py index bd8c3be271505..291c9d1ee7f0c 100644 --- a/pandas/io/sas/sasreader.py +++ b/pandas/io/sas/sasreader.py @@ -1,11 +1,16 @@ """ Read SAS sas7bdat or xport files. """ - from abc import ABCMeta, abstractmethod +from typing import TYPE_CHECKING, Optional, Union, overload + +from pandas._typing import FilePathOrBuffer, Label from pandas.io.common import stringify_path +if TYPE_CHECKING: + from pandas import DataFrame # noqa: F401 + # TODO(PY38): replace with Protocol in Python 3.8 class ReaderBase(metaclass=ABCMeta): @@ -22,14 +27,38 @@ def close(self): pass +@overload +def read_sas( + filepath_or_buffer: FilePathOrBuffer, + format: Optional[str] = ..., + index: Optional[Label] = ..., + encoding: Optional[str] = ..., + chunksize: int = ..., + iterator: bool = ..., +) -> ReaderBase: + ... + + +@overload +def read_sas( + filepath_or_buffer: FilePathOrBuffer, + format: Optional[str] = ..., + index: Optional[Label] = ..., + encoding: Optional[str] = ..., + chunksize: None = ..., + iterator: bool = ..., +) -> Union["DataFrame", ReaderBase]: + ... + + def read_sas( - filepath_or_buffer, - format=None, - index=None, - encoding=None, - chunksize=None, - iterator=False, -): + filepath_or_buffer: FilePathOrBuffer, + format: Optional[str] = None, + index: Optional[Label] = None, + encoding: Optional[str] = None, + chunksize: Optional[int] = None, + iterator: bool = False, +) -> Union["DataFrame", ReaderBase]: """ Read SAS files stored as either XPORT or SAS7BDAT format files.