TYPING: some type hints for pandas\io\common.py (pandas-dev#27598)

Author: simonjayhawkins

pandas/core/frame.py

@@ -732,7 +732,6 @@ def to_string(
 
         formatter = fmt.DataFrameFormatter(
             self,
-            buf=buf,
             columns=columns,
             col_space=col_space,
             na_rep=na_rep,
@@ -750,11 +749,7 @@ def to_string(
             decimal=decimal,
             line_width=line_width,
         )
-        formatter.to_string()
-
-        if buf is None:
-            result = formatter.buf.getvalue()
-            return result
+        return formatter.to_string(buf=buf)
 
     # ----------------------------------------------------------------------
 
@@ -2273,7 +2268,6 @@ def to_html(
 
         formatter = fmt.DataFrameFormatter(
             self,
-            buf=buf,
             columns=columns,
             col_space=col_space,
             na_rep=na_rep,
@@ -2294,10 +2288,9 @@ def to_html(
             render_links=render_links,
         )
         # TODO: a generic formatter wld b in DataFrameFormatter
-        formatter.to_html(classes=classes, notebook=notebook, border=border)
-
-        if buf is None:
-            return formatter.buf.getvalue()
+        return formatter.to_html(
+            buf=buf, classes=classes, notebook=notebook, border=border
+        )
 
     # ----------------------------------------------------------------------
 

pandas/core/generic.py

@@ -3018,7 +3018,6 @@ def to_latex(
 
         formatter = DataFrameFormatter(
             self,
-            buf=buf,
             columns=columns,
             col_space=col_space,
             na_rep=na_rep,
@@ -3032,7 +3031,8 @@ def to_latex(
             escape=escape,
             decimal=decimal,
         )
-        formatter.to_latex(
+        return formatter.to_latex(
+            buf=buf,
             column_format=column_format,
             longtable=longtable,
             encoding=encoding,
@@ -3041,9 +3041,6 @@ def to_latex(
             multirow=multirow,
         )
 
-        if buf is None:
-            return formatter.buf.getvalue()
-
     def to_csv(
         self,
         path_or_buf=None,

pandas/io/common.py

@@ -10,6 +10,7 @@
 import mmap
 import os
 import pathlib
+from typing import IO, AnyStr, BinaryIO, Optional, TextIO, Type
 from urllib.error import URLError  # noqa
 from urllib.parse import (  # noqa
     urlencode,
@@ -32,6 +33,8 @@
 
 from pandas.core.dtypes.common import is_file_like
 
+from pandas._typing import FilePathOrBuffer
+
 # gh-12665: Alias for now and remove later.
 CParserError = ParserError
 
@@ -68,14 +71,14 @@ class BaseIterator:
     Useful only when the object being iterated is non-reusable (e.g. OK for a
     parser, not for an in-memory table, yes for its iterator)."""
 
-    def __iter__(self):
+    def __iter__(self) -> "BaseIterator":
         return self
 
     def __next__(self):
         raise AbstractMethodError(self)
 
 
-def _is_url(url):
+def _is_url(url) -> bool:
     """Check to see if a URL has a valid protocol.
 
     Parameters
@@ -93,7 +96,9 @@ def _is_url(url):
         return False
 
 
-def _expand_user(filepath_or_buffer):
+def _expand_user(
+    filepath_or_buffer: FilePathOrBuffer[AnyStr]
+) -> FilePathOrBuffer[AnyStr]:
     """Return the argument with an initial component of ~ or ~user
        replaced by that user's home directory.
 
@@ -111,7 +116,7 @@ def _expand_user(filepath_or_buffer):
     return filepath_or_buffer
 
 
-def _validate_header_arg(header):
+def _validate_header_arg(header) -> None:
     if isinstance(header, bool):
         raise TypeError(
             "Passing a bool to header is invalid. "
@@ -121,7 +126,9 @@ def _validate_header_arg(header):
         )
 
 
-def _stringify_path(filepath_or_buffer):
+def _stringify_path(
+    filepath_or_buffer: FilePathOrBuffer[AnyStr]
+) -> FilePathOrBuffer[AnyStr]:
     """Attempt to convert a path-like object to a string.
 
     Parameters
@@ -144,21 +151,22 @@ def _stringify_path(filepath_or_buffer):
     strings, buffers, or anything else that's not even path-like.
     """
     if hasattr(filepath_or_buffer, "__fspath__"):
-        return filepath_or_buffer.__fspath__()
+        # https://github.com/python/mypy/issues/1424
+        return filepath_or_buffer.__fspath__()  # type: ignore
     elif isinstance(filepath_or_buffer, pathlib.Path):
         return str(filepath_or_buffer)
     return _expand_user(filepath_or_buffer)
 
 
-def is_s3_url(url):
+def is_s3_url(url) -> bool:
     """Check for an s3, s3n, or s3a url"""
     try:
         return parse_url(url).scheme in ["s3", "s3n", "s3a"]
     except Exception:
         return False
 
 
-def is_gcs_url(url):
+def is_gcs_url(url) -> bool:
     """Check for a gcs url"""
     try:
         return parse_url(url).scheme in ["gcs", "gs"]
@@ -167,7 +175,10 @@ def is_gcs_url(url):
 
 
 def get_filepath_or_buffer(
-    filepath_or_buffer, encoding=None, compression=None, mode=None
+    filepath_or_buffer: FilePathOrBuffer,
+    encoding: Optional[str] = None,
+    compression: Optional[str] = None,
+    mode: Optional[str] = None,
 ):
     """
     If the filepath_or_buffer is a url, translate and return the buffer.
@@ -190,7 +201,7 @@ def get_filepath_or_buffer(
     """
     filepath_or_buffer = _stringify_path(filepath_or_buffer)
 
-    if _is_url(filepath_or_buffer):
+    if isinstance(filepath_or_buffer, str) and _is_url(filepath_or_buffer):
         req = urlopen(filepath_or_buffer)
         content_encoding = req.headers.get("Content-Encoding", None)
         if content_encoding == "gzip":
@@ -224,7 +235,7 @@ def get_filepath_or_buffer(
     return filepath_or_buffer, None, compression, False
 
 
-def file_path_to_url(path):
+def file_path_to_url(path: str) -> str:
     """
     converts an absolute native path to a FILE URL.
 
@@ -242,7 +253,9 @@ def file_path_to_url(path):
 _compression_to_extension = {"gzip": ".gz", "bz2": ".bz2", "zip": ".zip", "xz": ".xz"}
 
 
-def _infer_compression(filepath_or_buffer, compression):
+def _infer_compression(
+    filepath_or_buffer: FilePathOrBuffer, compression: Optional[str]
+) -> Optional[str]:
     """
     Get the compression method for filepath_or_buffer. If compression='infer',
     the inferred compression method is returned. Otherwise, the input
@@ -435,7 +448,13 @@ class BytesZipFile(zipfile.ZipFile, BytesIO):  # type: ignore
     """
 
     # GH 17778
-    def __init__(self, file, mode, compression=zipfile.ZIP_DEFLATED, **kwargs):
+    def __init__(
+        self,
+        file: FilePathOrBuffer,
+        mode: str,
+        compression: int = zipfile.ZIP_DEFLATED,
+        **kwargs
+    ):
         if mode in ["wb", "rb"]:
             mode = mode.replace("b", "")
         super().__init__(file, mode, compression, **kwargs)
@@ -461,16 +480,16 @@ class MMapWrapper(BaseIterator):
 
     """
 
-    def __init__(self, f):
+    def __init__(self, f: IO):
         self.mmap = mmap.mmap(f.fileno(), 0, access=mmap.ACCESS_READ)
 
-    def __getattr__(self, name):
+    def __getattr__(self, name: str):
         return getattr(self.mmap, name)
 
-    def __iter__(self):
+    def __iter__(self) -> "MMapWrapper":
         return self
 
-    def __next__(self):
+    def __next__(self) -> str:
         newline = self.mmap.readline()
 
         # readline returns bytes, not str, but Python's CSV reader
@@ -491,16 +510,16 @@ class UTF8Recoder(BaseIterator):
     Iterator that reads an encoded stream and re-encodes the input to UTF-8
     """
 
-    def __init__(self, f, encoding):
+    def __init__(self, f: BinaryIO, encoding: str):
         self.reader = codecs.getreader(encoding)(f)
 
-    def read(self, bytes=-1):
+    def read(self, bytes: int = -1) -> bytes:
         return self.reader.read(bytes).encode("utf-8")
 
-    def readline(self):
+    def readline(self) -> bytes:
         return self.reader.readline().encode("utf-8")
 
-    def next(self):
+    def next(self) -> bytes:
         return next(self.reader).encode("utf-8")
 
 
@@ -511,5 +530,7 @@ def UnicodeReader(f, dialect=csv.excel, encoding="utf-8", **kwds):
     return csv.reader(f, dialect=dialect, **kwds)
 
 
-def UnicodeWriter(f, dialect=csv.excel, encoding="utf-8", **kwds):
+def UnicodeWriter(
+    f: TextIO, dialect: Type[csv.Dialect] = csv.excel, encoding: str = "utf-8", **kwds
+):
     return csv.writer(f, dialect=dialect, **kwds)