diff --git a/pdfplumber/display.py b/pdfplumber/display.py index c6e83f68..58a79903 100644 --- a/pdfplumber/display.py +++ b/pdfplumber/display.py @@ -33,7 +33,9 @@ class COLORS: T_contains_points = Union[Tuple[T_point, ...], List[T_point], T_obj] -def get_page_image(stream: BufferedReader, page_no: int, resolution: int) -> WandImage: +def get_page_image( + stream: Union[BufferedReader, BytesIO], page_no: int, resolution: int +) -> WandImage: # If we are working with a file object saved to disk if hasattr(stream, "name"): filename = f"{stream.name}[{page_no}]" diff --git a/pdfplumber/pdf.py b/pdfplumber/pdf.py index 4f16e494..3e152574 100644 --- a/pdfplumber/pdf.py +++ b/pdfplumber/pdf.py @@ -1,7 +1,7 @@ import itertools import logging import pathlib -from io import BufferedReader +from io import BufferedReader, BytesIO from types import TracebackType from typing import Any, Dict, List, Optional, Tuple, Type, Union @@ -25,7 +25,7 @@ class PDF(Container): def __init__( self, - stream: BufferedReader, + stream: Union[BufferedReader, BytesIO], stream_is_external: bool = False, pages: Optional[Union[List[int], Tuple[int]]] = None, laparams: Optional[Dict[str, Any]] = None, @@ -60,7 +60,7 @@ def __init__( @classmethod def open( cls, - path_or_fp: Union[str, pathlib.Path, BufferedReader], + path_or_fp: Union[str, pathlib.Path, BufferedReader, BytesIO], pages: Optional[Union[List[int], Tuple[int]]] = None, laparams: Optional[Dict[str, Any]] = None, password: str = "", @@ -68,7 +68,7 @@ def open( ) -> "PDF": if isinstance(path_or_fp, (str, pathlib.Path)): - stream = open(path_or_fp, "rb") + stream: Union[BufferedReader, BytesIO] = open(path_or_fp, "rb") stream_is_external = False else: stream = path_or_fp