Skip to content

Commit 44a828c

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 51fbfa3 commit 44a828c

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
@@ -35,6 +35,7 @@
3535
from datetime import datetime
3636
from io import BytesIO, FileIO, UnsupportedOperation
3737
from pathlib import Path
38+
from types import TracebackType
3839
from typing import (
3940
Any,
4041
Callable,
@@ -45,6 +46,7 @@
4546
Mapping,
4647
Optional,
4748
Tuple,
49+
Type,
4850
Union,
4951
cast,
5052
)
@@ -278,6 +280,9 @@ class PdfReader:
278280
password: Decrypt PDF file at initialization. If the
279281
password is None, the file will not be decrypted.
280282
Defaults to ``None``
283+
284+
Can also be instantiated as a contextmanager which will automatically close
285+
the underlying file pointer if passed via filenames.
281286
"""
282287

283288
@property
@@ -312,8 +317,10 @@ def __init__(
312317
__name__,
313318
)
314319

320+
self._opened_automatically = False
315321
if isinstance(stream, (str, Path)):
316322
stream = FileIO(stream, "rb")
323+
self._opened_automatically = True
317324
weakref.finalize(self, stream.close)
318325

319326
self.read(stream)
@@ -349,6 +356,20 @@ def close(self) -> None:
349356
"""Close the underlying file handle"""
350357
self.stream.close()
351358

359+
def __enter__(self) -> "PdfReader":
360+
"""Use PdfReader as context manager"""
361+
return self
362+
363+
def __exit__(
364+
self,
365+
exc_type: Optional[Type[BaseException]],
366+
exc: Optional[BaseException],
367+
traceback: Optional[TracebackType],
368+
) -> None:
369+
"""Close the underlying stream if owned by the PdfReader"""
370+
if self._opened_automatically:
371+
self.close()
372+
352373
@property
353374
def root_object(self) -> DictionaryObject:
354375
"""Provide access to "/Root". standardized with PdfWriter."""

tests/test_reader.py

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

850850
def test_read_path():
851851
path = Path(RESOURCE_ROOT, "crazyones.pdf")
852-
reader = PdfReader(path)
853-
assert len(reader.pages) == 1
852+
with PdfReader(path) as reader:
853+
assert len(reader.pages) == 1
854854

855855

856856
def test_read_not_binary_mode(caplog):

0 commit comments

Comments
 (0)