Skip to content

Commit 3578bf7

Browse files
committed
MAINT: Allow opening PdfReader as contextmanager
To mirror PdfWriter, also hints towards file pointer management now that we keep files open sometimes.
1 parent c8a77aa commit 3578bf7

File tree

2 files changed

+23
-2
lines changed

2 files changed

+23
-2
lines changed

pypdf/_reader.py

+21
Original file line numberDiff line numberDiff line change
@@ -32,6 +32,7 @@
3232
import weakref
3333
from io import BytesIO, FileIO, UnsupportedOperation
3434
from pathlib import Path
35+
from types import TracebackType
3536
from typing import (
3637
Any,
3738
Callable,
@@ -40,6 +41,7 @@
4041
List,
4142
Optional,
4243
Tuple,
44+
Type,
4345
Union,
4446
cast,
4547
)
@@ -100,6 +102,9 @@ class PdfReader(PdfDocCommon):
100102
password: Decrypt PDF file at initialization. If the
101103
password is None, the file will not be decrypted.
102104
Defaults to ``None``.
105+
106+
Can also be instantiated as a contextmanager which will automatically close
107+
the underlying file pointer if passed via filenames.
103108
"""
104109

105110
def __init__(
@@ -123,8 +128,10 @@ def __init__(
123128
__name__,
124129
)
125130

131+
self._opened_automatically = False
126132
if isinstance(stream, (str, Path)):
127133
stream = FileIO(stream, "rb")
134+
self._opened_automatically = True
128135
weakref.finalize(self, stream.close)
129136

130137
self.read(stream)
@@ -160,6 +167,20 @@ def close(self) -> None:
160167
"""Close the underlying file handle"""
161168
self.stream.close()
162169

170+
def __enter__(self) -> "PdfReader":
171+
"""Use PdfReader as context manager"""
172+
return self
173+
174+
def __exit__(
175+
self,
176+
exc_type: Optional[Type[BaseException]],
177+
exc: Optional[BaseException],
178+
traceback: Optional[TracebackType],
179+
) -> None:
180+
"""Close the underlying stream if owned by the PdfReader"""
181+
if self._opened_automatically:
182+
self.close()
183+
163184
@property
164185
def root_object(self) -> DictionaryObject:
165186
"""Provide access to "/Root". standardized with PdfWriter."""

tests/test_reader.py

+2-2
Original file line numberDiff line numberDiff line change
@@ -853,8 +853,8 @@ def test_extract_text_hello_world():
853853

854854
def test_read_path():
855855
path = Path(RESOURCE_ROOT, "crazyones.pdf")
856-
reader = PdfReader(path)
857-
assert len(reader.pages) == 1
856+
with PdfReader(path) as reader:
857+
assert len(reader.pages) == 1
858858

859859

860860
def test_read_not_binary_mode(caplog):

0 commit comments

Comments
 (0)