feat: add in a timeout option

Signed-off-by: Akhil Narang <[email protected]>

Author: akhilnarang

pdfkit/api.py

@@ -5,7 +5,7 @@
 
 
 def from_url(url, output_path=None, options=None, toc=None, cover=None,
-             configuration=None, cover_first=False, verbose=False):
+             configuration=None, cover_first=False, verbose=False, timeout=None):
     """
     Convert file of files from URLs to PDF document
 
@@ -17,18 +17,19 @@ def from_url(url, output_path=None, options=None, toc=None, cover=None,
     :param configuration: (optional) instance of pdfkit.configuration.Configuration()
     :param cover_first: (optional) if True, cover always precedes TOC
     :param verbose: (optional) By default '--quiet' is passed to all calls, set this to False to get wkhtmltopdf output to stdout.
+    :param timeout: (optional) if passed, uses this timeout value (in seconds) for the wkhtmltopdf command
 
     Returns: True on success
     """
 
     r = PDFKit(url, 'url', options=options, toc=toc, cover=cover,
-               configuration=configuration, cover_first=cover_first, verbose=verbose)
+               configuration=configuration, cover_first=cover_first, verbose=verbose, timeout=timeout)
 
     return r.to_pdf(output_path)
 
 
 def from_file(input, output_path=None, options=None, toc=None, cover=None, css=None,
-              configuration=None, cover_first=False, verbose=False):
+              configuration=None, cover_first=False, verbose=False, timeout=None):
     """
     Convert HTML file or files to PDF document
 
@@ -41,18 +42,19 @@ def from_file(input, output_path=None, options=None, toc=None, cover=None, css=N
     :param configuration: (optional) instance of pdfkit.configuration.Configuration()
     :param cover_first: (optional) if True, cover always precedes TOC
     :param verbose: (optional) By default '--quiet' is passed to all calls, set this to False to get wkhtmltopdf output to stdout.
+    :param timeout: (optional) if passed, uses this timeout value (in seconds) for the wkhtmltopdf command
 
     Returns: True on success
     """
 
     r = PDFKit(input, 'file', options=options, toc=toc, cover=cover, css=css,
-               configuration=configuration, cover_first=cover_first, verbose=verbose)
+               configuration=configuration, cover_first=cover_first, verbose=verbose, timeout=timeout)
 
     return r.to_pdf(output_path)
 
 
 def from_string(input, output_path=None, options=None, toc=None, cover=None, css=None,
-                configuration=None, cover_first=False, verbose=False):
+                configuration=None, cover_first=False, verbose=False, timeout=None):
     """
     Convert given string or strings to PDF document
 
@@ -65,12 +67,13 @@ def from_string(input, output_path=None, options=None, toc=None, cover=None, css
     :param configuration: (optional) instance of pdfkit.configuration.Configuration()
     :param cover_first: (optional) if True, cover always precedes TOC
     :param verbose: (optional) By default '--quiet' is passed to all calls, set this to False to get wkhtmltopdf output to stdout.
+    :param timeout: (optional) if passed, uses this timeout value (in seconds) for the wkhtmltopdf command
 
     Returns: True on success
     """
 
     r = PDFKit(input, 'string', options=options, toc=toc, cover=cover, css=css,
-               configuration=configuration, cover_first=cover_first, verbose=verbose)
+               configuration=configuration, cover_first=cover_first, verbose=verbose, timeout=timeout)
 
     return r.to_pdf(output_path)
 

pdfkit/pdfkit.py

@@ -32,7 +32,7 @@ def __str__(self):
             return self.msg
 
     def __init__(self, url_or_file, type_, options=None, toc=None, cover=None,
-                 css=None, configuration=None, cover_first=False, verbose=False):
+                 css=None, configuration=None, cover_first=False, verbose=False, timeout=None):
 
         self.source = Source(url_or_file, type_)
         self.configuration = (Configuration() if configuration is None
@@ -57,6 +57,7 @@ def __init__(self, url_or_file, type_, options=None, toc=None, cover=None,
         self.verbose = verbose
         self.css = css
         self.stylesheets = []
+        self.timeout = timeout
 
     def _genargs(self, opts):
         """
@@ -187,7 +188,13 @@ def to_pdf(self, path=None):
         else:
             input = None
 
-        stdout, stderr = result.communicate(input=input)
+        kwargs = {"input": input}
+
+        # Check if a timeout was specified and pass it along if so
+        if self.timeout:
+            kwargs["timeout"] = self.timeout
+
+        stdout, stderr = result.communicate(**kwargs)
         stderr = stderr or stdout or b""
         stderr = stderr.decode('utf-8', errors='replace')
         exit_code = result.returncode