Adding Timeout CLI parameter

pdfx/cli.py

@@ -94,6 +94,13 @@ def create_parser():
         action="version",
         version="%(prog)s v{version}".format(version=pdfx.__version__),
     )
+
+    parser.add_argument(
+        "-f",
+        "--fast-timeout",
+        action="store_true",
+        help="Apply 10 second timeout for URLOPEN to prevent hanging processes",
+    )
     return parser
 
 
@@ -201,7 +208,10 @@ def main():
         refs_all = pdf.get_references()
         refs = [ref for ref in refs_all if ref.reftype in ["url", "pdf"]]
         print("\nChecking %s URLs for broken links..." % len(refs))
-        check_refs(refs)
+        if args.fast_timeout:
+            check_refs(refs, timeout=True)
+        else:
+            check_refs(refs)
 
     try:
         if args.download_pdfs:

pdfx/downloader.py

@@ -3,6 +3,7 @@
 from .colorprint import colorprint, OKGREEN, FAIL
 from .threadpool import ThreadPool
 from collections import defaultdict
+from itertools import repeat
 import ssl
 import os
 import sys
@@ -37,7 +38,7 @@ def sanitize_url(url):
     return url
 
 
-def get_status_code(url):
+def get_status_code(url, timeout):
     """ Perform HEAD request and return status code """
     try:
         request = Request(sanitize_url(url))
@@ -46,7 +47,10 @@ def get_status_code(url):
             "Mozilla/5.0 (compatible; MSIE 9.0; " "Windows NT 6.1; Trident/5.0)",
         )
         request.get_method = lambda: "HEAD"
-        response = urlopen(request, context=ssl_unverified_context)
+        if timeout:
+            response = urlopen(request, context=ssl_unverified_context, timeout=10)
+        else:
+            response = urlopen(request, context=ssl_unverified_context)
         # print response.info()
         return response.getcode()
     except HTTPError as e:
@@ -58,13 +62,13 @@ def get_status_code(url):
         return None
 
 
-def check_refs(refs, verbose=True, max_threads=MAX_THREADS_DEFAULT):
+def check_refs(refs, verbose=True, max_threads=MAX_THREADS_DEFAULT, timeout=False):
     """ Check if urls exist """
     codes = defaultdict(list)
 
-    def check_url(ref):
+    def check_url(ref, timeout):
         url = ref.ref
-        status_code = str(get_status_code(url))
+        status_code = str(get_status_code(url, timeout))
         codes[status_code].append(ref)
         if verbose:
             if status_code == "200":
@@ -75,7 +79,7 @@ def check_url(ref):
     # Start a threadpool and add the check-url tasks
     try:
         pool = ThreadPool(5)
-        pool.map(check_url, refs)
+        pool.starmap(check_url, zip(refs, repeat(timeout)))
         pool.wait_completion()
 
     except Exception as e: