Allow to save result to file-like object. (#66)

Author: vb64

README.md

@@ -109,6 +109,16 @@ Save to file.
 pdf.save("guide.pdf")
 ```
 
+Or save to file-like object.
+
+```python
+import io
+
+out = io.BytesIO()
+pdf.save_bytes(out)
+assert out.getbuffer().nbytes > 0
+```
+
 ![Pdf](img/with_toc.png)
 
 ## Settings and options

README_ru.md

@@ -103,6 +103,16 @@ pdf.meta["author"] = "Виталий Богомолов"
 pdf.save("guide.pdf")
 ```
 
+Либо сохраняем в буфер памяти.
+
+```python
+import io
+
+out = io.BytesIO()
+pdf.save_bytes(out)
+assert out.getbuffer().nbytes > 0
+```
+
 ![Pdf](img/with_toc.png)
 
 ## Настройки и опции

markdown_pdf/__init__.py

@@ -3,9 +3,11 @@
 import typing
 import pathlib
 
-
-from markdown_it import MarkdownIt
 import fitz
+from markdown_it import MarkdownIt
+from pymupdf import (
+  _as_pdf_document, mupdf, JM_embedded_clean, JM_ensure_identity, JM_new_output_fileptr, ASSERT_PDF,
+)
 
 
 class Section:
@@ -94,16 +96,57 @@ def add_section(self, section: Section, user_css: typing.Optional[str] = None) -
 
         return html
 
-    def save(self, file_name: typing.Union[str, pathlib.Path]) -> None:
-        """Save pdf to file."""
+    def _make_doc(self):
+        """Return fitz doc object."""
         self.writer.close()
         self.out_file.seek(0)
         doc = fitz.Story.add_pdf_links(self.out_file, self.hrefs)
         doc.set_metadata(self.meta)
         if self.toc_level > 0:
             doc.set_toc(self.toc)
+
+        return doc
+
+    def save(self, file_name: typing.Union[str, pathlib.Path]) -> None:
+        """Save pdf to file."""
+        doc = self._make_doc()
         if self.optimize:
             doc.ez_save(file_name)
         else:
             doc.save(file_name)
         doc.close()
+
+    def save_bytes(self, bytesio: io.BytesIO) -> None:
+        """Save pdf to file-like object."""
+        doc = self._make_doc()
+
+        pdf = _as_pdf_document(doc)
+        opts = mupdf.PdfWriteOptions()
+
+        opts.do_incremental = 0
+        opts.do_ascii = 0
+        opts.do_compress = 0
+        opts.do_compress_images = 0
+        opts.do_compress_fonts = 0
+        opts.do_decompress = 0
+        opts.do_garbage = 0
+        opts.do_pretty = 0
+        opts.do_linear = 0
+        opts.do_clean = 0
+        opts.do_sanitize = 0
+        opts.dont_regenerate_id = 0
+        opts.do_appearance = 0
+        opts.do_encrypt = 1
+        opts.permissions = 4095
+        opts.do_preserve_metadata = 1
+        opts.do_use_objstms = 0
+        opts.compression_effort = 0
+
+        ASSERT_PDF(pdf)
+        pdf.m_internal.resynth_required = 0
+        JM_embedded_clean(pdf)
+        JM_ensure_identity(pdf)
+
+        out = JM_new_output_fileptr(bytesio)
+        mupdf.pdf_write_document(pdf, out, opts)
+        out.fz_close_output()

tests/test/test_converter.py

@@ -2,6 +2,7 @@
 
 make test T=test_converter.py
 """
+import io
 from . import TestBase
 
 TABLE_TEXT = """# Section with Table
@@ -80,3 +81,16 @@ def test_hrefs(self):
           Section(open(self.fixture("hrefs.md"), "rt", encoding='utf-8').read())
         )
         pdf.save(self.build("hrefs.pdf"))
+
+    def test_bytes(self):
+        """Save pdf as bytes IO."""
+        from markdown_pdf import Section, MarkdownPdf
+        out = io.BytesIO()
+
+        pdf = MarkdownPdf(toc_level=2)
+        pdf.add_section(Section("# Title Bytes\n", toc=False))
+        pdf.add_section(Section("# Bytes1\n\nbody\n"))
+        pdf.save_bytes(out)
+        assert out.getbuffer().nbytes > 0
+        with open(self.build("as_bytes.pdf"), "wb") as i:
+            i.write(out.getvalue())