Merge pull request matplotlib#25655 from ksunden/file_io_cleanup

Clean up FileIO type hints

Author: QuLogic

lib/matplotlib/backend_bases.pyi

@@ -1,6 +1,5 @@
 from enum import Enum, IntEnum
 import os
-import io
 from matplotlib import (
     cbook,
     colors,
@@ -27,7 +26,7 @@ from matplotlib.text import Text
 from matplotlib.transforms import Affine2D, Transform, TransformedPath, Bbox
 
 from collections.abc import Callable, Iterable, Sequence
-from typing import Any, Literal, NamedTuple, TypeVar
+from typing import Any, IO, Literal, NamedTuple, TypeVar
 from numpy.typing import ArrayLike
 from .typing import ColorType, LineStyleType
 
@@ -347,7 +346,7 @@ class FigureCanvasBase:
     def get_supported_filetypes_grouped(cls) -> dict[str, list[str]]: ...
     def print_figure(
         self,
-        filename: str | os.PathLike | io.FileIO,
+        filename: str | os.PathLike | IO,
         dpi: float | None = ...,
         facecolor: ColorType | Literal["auto"] | None = ...,
         edgecolor: ColorType | Literal["auto"] | None = ...,

lib/matplotlib/cbook.pyi

@@ -1,7 +1,6 @@
 import collections.abc
 from collections.abc import Callable, Collection, Generator, Iterable, Iterator
 import contextlib
-import io
 import os
 from pathlib import Path
 
@@ -12,7 +11,9 @@ from numpy.typing import ArrayLike
 
 from typing import (
     Any,
+    ContextManager,
     Generic,
+    IO,
     Literal,
     TypeVar,
     overload,
@@ -42,35 +43,50 @@ class silent_list(list[_T]):
 def strip_math(s: str) -> str: ...
 def is_writable_file_like(obj: Any) -> bool: ...
 def file_requires_unicode(x: Any) -> bool: ...
+@overload
 def to_filehandle(
-    fname: str | os.PathLike | io.FileIO,
+    fname: str | os.PathLike | IO,
     flag: str = ...,
-    return_opened: bool = ...,
+    return_opened: Literal[False] = ...,
+    encoding: str | None = ...,
+) -> IO: ...
+@overload
+def to_filehandle(
+    fname: str | os.PathLike | IO,
+    flag: str,
+    return_opened: Literal[True],
+    encoding: str | None = ...,
+) -> tuple[IO, bool]: ...
+@overload
+def to_filehandle(
+    fname: str | os.PathLike | IO,
+    *, # if flag given, will match previous sig
+    return_opened: Literal[True],
     encoding: str | None = ...,
-) -> io.FileIO: ...
+) -> tuple[IO, bool]: ...
 def open_file_cm(
-    path_or_file: str | os.PathLike | io.FileIO,
+    path_or_file: str | os.PathLike | IO,
     mode: str = ...,
     encoding: str | None = ...,
-): ...
+) -> ContextManager[IO]: ...
 def is_scalar_or_string(val: Any) -> bool: ...
 @overload
 def get_sample_data(
-    fname: str | os.PathLike | io.FileIO,
+    fname: str | os.PathLike,
     asfileobj: Literal[True] = ...,
     *,
     np_load: Literal[True]
 ) -> np.ndarray: ...
 @overload
 def get_sample_data(
-    fname: str | os.PathLike | io.FileIO,
+    fname: str | os.PathLike,
     asfileobj: Literal[True] = ...,
     *,
     np_load: Literal[False] = ...
-) -> io.FileIO: ...
+) -> IO: ...
 @overload
 def get_sample_data(
-    fname: str | os.PathLike | io.FileIO,
+    fname: str | os.PathLike,
     asfileobj: Literal[False],
     *,
     np_load: bool = ...
@@ -96,7 +112,7 @@ class Stack(Generic[_T]):
 
 def safe_masked_invalid(x: ArrayLike, copy: bool = ...) -> np.ndarray: ...
 def print_cycles(
-    objects: Iterable[Any], outstream: io.FileIO = ..., show_progress: bool = ...
+    objects: Iterable[Any], outstream: IO = ..., show_progress: bool = ...
 ) -> None: ...
 
 class Grouper(Generic[_T]):

lib/matplotlib/figure.pyi

@@ -1,4 +1,3 @@
-import io
 import os
 
 from matplotlib import backend_bases, projections
@@ -40,7 +39,7 @@ import numpy as np
 from numpy.typing import ArrayLike
 
 from collections.abc import Callable, Iterable
-from typing import Any, Literal, overload
+from typing import Any, IO, Literal, overload
 from .typing import ColorType, HashableList
 
 class SubplotParams:
@@ -372,7 +371,7 @@ class Figure(FigureBase):
     def add_axobserver(self, func: Callable[[Figure], Any]): ...
     def savefig(
         self,
-        fname: str | os.PathLike | io.FileIO,
+        fname: str | os.PathLike | IO,
         *,
         transparent: bool | None = ...,
         **kwargs

lib/matplotlib/image.pyi

@@ -21,7 +21,7 @@ from matplotlib.transforms import (
 )
 
 from collections.abc import Sequence
-from typing import Any, Literal
+from typing import Any, BinaryIO, Literal
 import numpy as np
 from numpy.typing import ArrayLike
 
@@ -76,7 +76,7 @@ class _ImageBase(martist.Artist, cm.ScalarMappable):
         self, renderer: RendererBase, magnification: float = ..., unsampled: bool = ...
     ) -> tuple[np.ndarray, float, float, Affine2D]: ...
     def draw(self, renderer: RendererBase, *args, **kwargs) -> None: ...
-    def write_png(self, fname: str | pathlib.Path | io.FileIO) -> None: ...
+    def write_png(self, fname: str | pathlib.Path | BinaryIO) -> None: ...
     def set_data(self, A: ArrayLike | None) -> None: ...
     def set_array(self, A: ArrayLike | None) -> None: ...
     def get_shape(self) -> tuple[int, int, int]: ...
@@ -176,9 +176,11 @@ class BboxImage(_ImageBase):
     ) -> None: ...
     def get_window_extent(self, renderer: RendererBase | None = ...): ...
 
-def imread(fname: str | io.FileIO, format: str | None = ...) -> np.ndarray: ...
+def imread(
+    fname: str | pathlib.Path |  BinaryIO, format: str | None = ...
+) -> np.ndarray: ...
 def imsave(
-    fname: str | os.PathLike | io.FileIO,
+    fname: str | os.PathLike | BinaryIO,
     arr: ArrayLike,
     vmin: float | None = ...,
     vmax: float | None = ...,
@@ -192,8 +194,8 @@ def imsave(
 ) -> None: ...
 def pil_to_array(pilImage: PIL.Image.Image) -> np.ndarray: ...
 def thumbnail(
-    infile: str | io.FileIO,
-    thumbfile: str | io.FileIO,
+    infile: str | BinaryIO,
+    thumbfile: str | BinaryIO,
     scale: float = ...,
     interpolation: str = ...,
     preview: bool = ...,

lib/matplotlib/mathtext.pyi

@@ -4,7 +4,7 @@ from ._mathtext import RasterParse, VectorParse, get_unicode_index
 from matplotlib.font_manager import FontProperties
 from matplotlib.ft2font import FT2Image, LOAD_NO_HINTING
 
-from typing import Literal
+from typing import IO, Literal
 from matplotlib.typing import ColorType
 
 class MathTextParser:
@@ -13,7 +13,7 @@ class MathTextParser:
 
 def math_to_image(
     s: str,
-    filename_or_obj: str | os.PathLike | io.FileIO,
+    filename_or_obj: str | os.PathLike | IO,
     prop: FontProperties | None = ...,
     dpi: float | None = ...,
     format: str | None = ...,

lib/matplotlib/pyplot.py

@@ -82,15 +82,14 @@
 
 if TYPE_CHECKING:
     import datetime
-    import io
     import pathlib
     import os
 
     import PIL
 
     from numpy.typing import ArrayLike
     from typing import (
-        Any, Callable, Hashable, Literal, Sequence, Iterable, Type, Generator
+        Any, BinaryIO, Callable, Hashable, Literal, Sequence, Iterable, Type
     )
 
     from matplotlib.axis import Tick
@@ -2289,13 +2288,15 @@ def set_cmap(cmap: Colormap | str) -> None:
 
 
 @_copy_docstring_and_deprecators(matplotlib.image.imread)
-def imread(fname: str | io.FileIO, format: str | None = None) -> np.ndarray:
+def imread(
+        fname: str | pathlib.Path | BinaryIO, format: str | None = None
+) -> np.ndarray:
     return matplotlib.image.imread(fname, format)
 
 
 @_copy_docstring_and_deprecators(matplotlib.image.imsave)
 def imsave(
-    fname: str | os.PathLike | io.FileIO, arr: ArrayLike, **kwargs
+    fname: str | os.PathLike | BinaryIO, arr: ArrayLike, **kwargs
 ) -> None:
     return matplotlib.image.imsave(fname, arr, **kwargs)