Merge branch 'Kazy-pikepdf'

6 files changed, 59 insertions, 31 deletions

diff --git a/paper2remarkable/crop.py b/paper2remarkable/crop.py
index 623d29f..16d050e 100644
--- a/paper2remarkable/crop.py
+++ b/paper2remarkable/crop.py
@@ -8,13 +8,12 @@ Copyright: 2019, G.J.J. van den Burg
 
 """
 
-import PyPDF2
 import io
 import os
 import pdfplumber
 import subprocess
 
-from PyPDF2.generic import RectangleObject
+from pikepdf import Pdf
 
 from .log import Logger
 
@@ -64,7 +63,7 @@ class Cropper(object):
     ):
         if not input_file is None:
             self.input_file = os.path.abspath(input_file)
-            self.reader = PyPDF2.PdfFileReader(self.input_file)
+            self.reader = Pdf.open(self.input_file)
         if not output_file is None:
             self.output_file = os.path.abspath(output_file)
 
@@ -72,7 +71,6 @@ class Cropper(object):
             pdftoppm_path = None
 
         self.pdftoppm_path = pdftoppm_path
-        self.writer = PyPDF2.PdfFileWriter()
 
     def crop(self, margins=1):
         return self.process_file(self.crop_page, margins=margins)
@@ -84,15 +82,16 @@ class Cropper(object):
         return self.process_file(self.right_page, padding=padding)
 
     def process_file(self, page_func, *args, **kwargs):
-        n = self.reader.getNumPages()
+        n = len(self.reader.pages)
         for page_idx in range(n):
             status = page_func(page_idx, *args, **kwargs)
             if not status == 0:
                 return status
             if (page_idx + 1) % 10 == 0:
                 logger.info("Processing pages ... (%i/%i)" % (page_idx + 1, n))
-        with open(self.output_file, "wb") as fp:
-            self.writer.write(fp)
+
+        self.reader.save(self.output_file)
+        self.reader.close()
         if n % 10 > 0:
             logger.info("Processing pages ... (%i/%i)" % (n, n))
         return 0
@@ -112,21 +111,25 @@ class Cropper(object):
 
     def export_page(self, page_idx):
         """Helper function that exports a single page given by index """
-        page = self.reader.getPage(page_idx)
-        writer = PyPDF2.PdfFileWriter()
-        writer.addPage(page)
+        page = self.reader.pages[page_idx]
+
+        writer = Pdf.new()
+        writer.pages.append(page)
+        # Remove the annotations to avoid warning about `Bad annotation destination` when
+        # processing the page with pdftoppm. Since we've appended the page to the writer, pikepdf
+        # does a copy of the copy and we're not modifying the original page's annotations.
+        writer.pages[0].Annots = []
+
         tmpfname = "./page.pdf"
-        with open(tmpfname, "wb") as fp:
-            writer.write(fp)
+        writer.save(tmpfname)
+        writer.close()
         return tmpfname
 
     def process_page(self, page_idx, bbox_func, *args, **kwargs):
         """Process a single page and add it to the writer """
         tmpfname = self.export_page(page_idx)
         bbox = bbox_func(tmpfname, *args, **kwargs)
-        thepage = self.reader.getPage(page_idx)
-        thepage.cropBox = RectangleObject(bbox)
-        self.writer.addPage(thepage)
+        self.reader.pages[page_idx].CropBox = bbox
         os.unlink(tmpfname)
         return 0
 
diff --git a/paper2remarkable/pdf_ops.py b/paper2remarkable/pdf_ops.py
index c365920..93f1200 100644
--- a/paper2remarkable/pdf_ops.py
+++ b/paper2remarkable/pdf_ops.py
@@ -9,10 +9,11 @@ Copyright: 2019, G.J.J. van den Burg
 """
 
 
-import PyPDF2
 import os
 import subprocess
 
+from pikepdf import Pdf
+
 from .crop import Cropper
 from .log import Logger
 
@@ -42,15 +43,17 @@ def prepare_pdf(filepath, operation, pdftoppm_path="pdftoppm"):
 def blank_pdf(filepath):
     """Add blank pages to PDF"""
     logger.info("Adding blank pages")
-    input_pdf = PyPDF2.PdfFileReader(filepath)
-    output_pdf = PyPDF2.PdfFileWriter()
-    for page in input_pdf.pages:
-        output_pdf.addPage(page)
-        output_pdf.addBlankPage()
+    pdf = Pdf.open(filepath)
+
+    previous_pages = pdf.pages
+    pdf.pages = []
+
+    for page in previous_pages:
+        pdf.pages.append(page)
+        pdf.add_blank_page()
 
     output_file = os.path.splitext(filepath)[0] + "-blank.pdf"
-    with open(output_file, "wb") as fp:
-        output_pdf.write(fp)
+    pdf.save(output_file)
     return output_file
 
 
diff --git a/paper2remarkable/providers/_base.py b/paper2remarkable/providers/_base.py
index 0453c7a..369d566 100644
--- a/paper2remarkable/providers/_base.py
+++ b/paper2remarkable/providers/_base.py
@@ -86,6 +86,7 @@ class Provider(metaclass=abc.ABCMeta):
 
         if blank:
             self.operations.append(("blank", blank_pdf))
+
         self.operations.append(("shrink", self.shrink_pdf))
 
         logger.info("Starting %s provider" % type(self).__name__)
@@ -218,6 +219,7 @@ class Provider(metaclass=abc.ABCMeta):
             intermediate_fname = tmp_filename
             for opname, op in self.operations:
                 intermediate_fname = op(intermediate_fname)
+
             shutil.copy(intermediate_fname, clean_filename)
 
             if self.debug:
diff --git a/paper2remarkable/utils.py b/paper2remarkable/utils.py
index 2432916..0003103 100644
--- a/paper2remarkable/utils.py
+++ b/paper2remarkable/utils.py
@@ -8,7 +8,6 @@ Copyright: 2019, G.J.J. van den Burg
 
 """
 
-import PyPDF2
 import regex
 import requests
 import string
@@ -16,6 +15,8 @@ import subprocess
 import time
 import unidecode
 
+from pikepdf import Pdf, PdfError
+
 from .log import Logger
 from .exceptions import FileTypeError, RemarkableError, NoPDFToolError
 
@@ -45,15 +46,14 @@ def clean_string(s):
 def assert_file_is_pdf(filename):
     """Assert that a given file is a PDF file.
 
-    This is done by trying to open it using PyPDF2.
+    This is done by trying to open it using pikepdf.
     """
     try:
-        fp = open(filename, "rb")
-        pdf = PyPDF2.PdfFileReader(fp, strict=False)
-        fp.close()
+        pdf = Pdf.open(filename)
+        pdf.close()
         del pdf
         return True
-    except PyPDF2.utils.PdfReadError:
+    except PdfError:
         raise FileTypeError(filename, "pdf")
 
 
diff --git a/setup.py b/setup.py
index e529cc2..fb7d21c 100644
--- a/setup.py
+++ b/setup.py
@@ -13,13 +13,13 @@ EMAIL = "gertjanvandenburg@gmail.com"
 LICENSE = "MIT"
 LICENSE_TROVE = "License :: OSI Approved :: MIT License"
 NAME = "paper2remarkable"
-REQUIRES_PYTHON = ">=3.5.0"
+REQUIRES_PYTHON = ">=3.6.0"
 URL = "https://github.com/GjjvdBurg/paper2remarkable"
 VERSION = None
 
 # What packages are required for this module to be executed?
 REQUIRED = [
-    "PyPDF2>=1.26",
+    "pikepdf>=2.8.0",
     "beautifulsoup4>=4.8",
     "html2text>=2020.1.16",
     "markdown>=3.1.1",
diff --git a/tests/test_providers.py b/tests/test_providers.py
index eaeb8aa..af69c64 100644
--- a/tests/test_providers.py
+++ b/tests/test_providers.py
@@ -11,6 +11,7 @@ import pdfplumber
 import shutil
 import tempfile
 import unittest
+from pikepdf import Pdf
 
 from paper2remarkable.providers import (
     ACL,
@@ -34,6 +35,7 @@ from paper2remarkable.providers import (
     Springer,
     TandFOnline,
 )
+from paper2remarkable.utils import download_url
 
 VERBOSE = False
 
@@ -438,6 +440,24 @@ class TestProviders(unittest.TestCase):
         filename = prov.run(url)
         self.assertEqual(exp, os.path.basename(filename))
 
+    def test_local_file_copy_toc(self):
+        """Make sure the table of content is kept after processing."""
+        local_filename = "test.pdf"
+        download_url("https://arxiv.org/pdf/1711.03512.pdf", local_filename)
+        prov = LocalFile(upload=False, verbose=VERBOSE)
+        filename = prov.run(local_filename)
+        with Pdf.open(filename) as pdf:
+            with pdf.open_outline() as outline:
+                assert len(outline.root) > 0
+
+    def test_arxiv_copy_toc(self):
+        """Make sure the table of content is kept after processing when using the arXiv provider."""
+        prov = Arxiv(upload=False, verbose=VERBOSE)
+        filename = prov.run("https://arxiv.org/abs/1711.03512")
+        with Pdf.open(filename) as pdf:
+            with pdf.open_outline() as outline:
+                assert len(outline.root) > 0
+
 
 if __name__ == "__main__":
     unittest.main()