Use local DBLP XML dump instead of per-entry live fetches

.gitignore

@@ -12,3 +12,5 @@
 
 .DS_Store
 /auto
+/dblp-dump/
+/.venv/

README.md

@@ -2,4 +2,4 @@
 
 **WARNING**: This is probably not the repository your are interested in. This repository is only for *cryptobib* developers. The repositories containing the public *bib* files are [cryptobib/export](https://github.com/cryptobib/export) and  [cryptobib/export_crossref](https://github.com/cryptobib/export_crossref).
 
-**WARNING**: This project shall only be used as a subfolder of the main project [cryptobib/cryptobib](https://github.com/cryptobib/cryptobib). Please read the documentation of the main project.
+**WARNING**: This project shall only be used as a subfolder of the main project [cryptobib/cryptobib](https://github.com/cryptobib/cryptobib). Please read the documentation of the main project, including the setup instructions for `fetch_dblp_dump.py`.

fetch_dblp_dump.py

@@ -0,0 +1,116 @@
+#!/usr/bin/env python3
+# -*- coding: utf-8 -*-
+"""Fetch (and keep up to date) a local copy of the DBLP XML dump.
+
+import.py looks up publications in this local dump instead of hitting
+the DBLP website once per publication, which is what used to get the
+import script rate-limited/blocked on a heavily loaded DBLP.
+
+Usage:
+    python3 fetch_dblp_dump.py [--force] [--dir DIR]
+
+See: https://dblp.org/faq/How+can+I+download+the+whole+dblp+dataset.html
+"""
+
+import argparse
+import hashlib
+import os
+import sys
+import urllib.request
+
+DBLP_XML_URL = "https://dblp.org/xml/dblp.xml.gz"
+DBLP_XML_MD5_URL = "https://dblp.org/xml/dblp.xml.gz.md5"
+DBLP_DTD_URL = "https://dblp.org/xml/dblp.dtd"
+
+USER_AGENT = "cryptobib import script 1.0"
+
+scriptdir = os.path.dirname(os.path.realpath(__file__))
+DEFAULT_DUMP_DIR = os.path.join(scriptdir, "dblp-dump")
+
+
+def fetch(url):
+    req = urllib.request.Request(url, headers={"User-Agent": USER_AGENT})
+    with urllib.request.urlopen(req) as f:
+        return f.read()
+
+
+def fetch_to_file(url, path):
+    req = urllib.request.Request(url, headers={"User-Agent": USER_AGENT})
+    tmp_path = path + ".part"
+    with urllib.request.urlopen(req) as r, open(tmp_path, "wb") as out:
+        total = int(r.headers.get("Content-Length", 0))
+        read = 0
+        chunk_size = 1024 * 1024
+        while True:
+            chunk = r.read(chunk_size)
+            if not chunk:
+                break
+            out.write(chunk)
+            read += len(chunk)
+            if total:
+                print(
+                    "\r  {} / {} MiB ({:.0%})".format(
+                        read // (1024 * 1024), total // (1024 * 1024), read / total
+                    ),
+                    end="",
+                    file=sys.stderr,
+                )
+        print(file=sys.stderr)
+    os.replace(tmp_path, path)
+
+
+def md5sum(path):
+    h = hashlib.md5()
+    with open(path, "rb") as f:
+        for chunk in iter(lambda: f.read(1024 * 1024), b""):
+            h.update(chunk)
+    return h.hexdigest()
+
+
+def remote_md5():
+    # format: "<md5>  dblp.xml.gz"
+    return fetch(DBLP_XML_MD5_URL).decode("ascii").split()[0]
+
+
+def main():
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument(
+        "--dir",
+        default=DEFAULT_DUMP_DIR,
+        help="directory to store the dump in (default: {})".format(DEFAULT_DUMP_DIR),
+    )
+    parser.add_argument(
+        "--force", action="store_true", help="redownload even if local copy looks current"
+    )
+    args = parser.parse_args()
+
+    os.makedirs(args.dir, exist_ok=True)
+    gz_path = os.path.join(args.dir, "dblp.xml.gz")
+    dtd_path = os.path.join(args.dir, "dblp.dtd")
+
+    print("Checking remote dblp.xml.gz checksum...")
+    want_md5 = remote_md5()
+
+    if not args.force and os.path.exists(gz_path) and md5sum(gz_path) == want_md5:
+        print("Local dump already up to date ({}).".format(gz_path))
+    else:
+        print("Downloading {} -> {}".format(DBLP_XML_URL, gz_path))
+        fetch_to_file(DBLP_XML_URL, gz_path)
+        got_md5 = md5sum(gz_path)
+        if got_md5 != want_md5:
+            os.remove(gz_path)
+            sys.exit(
+                "Checksum mismatch for dblp.xml.gz (expected {}, got {}); download removed, try again.".format(
+                    want_md5, got_md5
+                )
+            )
+        print("Checksum OK.")
+
+    print("Downloading {} -> {}".format(DBLP_DTD_URL, dtd_path))
+    fetch_to_file(DBLP_DTD_URL, dtd_path)
+
+    print("Done. Dump ready in {}".format(args.dir))
+
+
+if __name__ == "__main__":
+    main()

import.py

@@ -28,9 +28,21 @@
 import argparse
 import html.parser
 import http
+import gzip
 
 from config import *
 
+# Local DBLP XML dump (see fetch_dblp_dump.py), used instead of hitting
+# dblp.org once per publication.
+DBLP_DUMP_DIR = os.environ.get("DBLP_DUMP_DIR", os.path.join(scriptdir, "dblp-dump"))
+DBLP_DUMP_GZ = os.path.join(DBLP_DUMP_DIR, "dblp.xml.gz")
+DBLP_DUMP_DTD = os.path.join(DBLP_DUMP_DIR, "dblp.dtd")
+
+DBLP_RECORD_TAGS = {
+    "article", "inproceedings", "proceedings", "book", "incollection",
+    "phdthesis", "mastersthesis", "www", "data",
+}
+
 logging_colorer.init()
 logging.basicConfig(level=logging.DEBUG)
 
@@ -959,19 +971,55 @@ def xml_get_value(e):
     r"^([0-9:]*)(--?([0-9:]*))?$"
 )  # LIPIcs uses pages of the form "5:1-5:10"
 
+re_dtd_entity = re.compile(r'<!ENTITY\s+(\w+)\s+"&#(\d+);"\s*>')
+
+
+def load_dtd_entities(dtd_path):
+    """parse the numeric-character-reference entities dblp.dtd defines (e.g. &Aacute; -> chr(193))"""
+    with open(dtd_path, encoding="ascii") as f:
+        return {m.group(1): chr(int(m.group(2))) for m in re_dtd_entity.finditer(f.read())}
+
+
+def lookup_dblp_records(gz_path, dtd_entities, wanted_keys):
+    """stream the local dblp.xml.gz dump once, returning {key: Element} for the requested keys found"""
+    wanted = set(wanted_keys)
+    found = {}
+    if not wanted:
+        return found
+    parser = xml.etree.ElementTree.XMLParser()
+    parser.entity.update(dtd_entities)
+    with gzip.open(gz_path, "rb") as f:
+        context = ElementTree.iterparse(f, events=("start", "end"), parser=parser)
+        _, root = next(context)
+        n = 0
+        for event, elem in context:
+            if event != "end" or elem.tag not in DBLP_RECORD_TAGS:
+                continue
+            key = elem.get("key")
+            if key in wanted:
+                found[key] = elem
+                wanted.discard(key)
+                if not wanted:
+                    break
+            else:
+                elem.clear()
+            n += 1
+            if n % 500000 == 0:
+                root.clear()
+                logging.info(
+                    "Scanned {} dblp dump records, {} key(s) left to find...".format(n, len(wanted))
+                )
+    if wanted:
+        logging.warning(
+            "{} key(s) not found in local dblp dump (maybe too recent, falling back to live fetch): {}".format(
+                len(wanted), ", ".join(sorted(wanted))
+            )
+        )
+    return found
 
-def xml_to_entry(xml, confkey, entry_type, fields, short_year, use_ee_as_url):
-    """transform a DBLP xml entry of type "entry_type" into a dictionnary ready to be output as bibtex"""
-    try:
-        tree = XML(xml)
-    except ElementTree.ParseError as e:
-        logging.exception("XML Parsing Error")
-        return None, None
-    elt = tree.find(entry_type.lower())
-    if elt is None:
-        logging.warning('Entry type is not "{0}"'.format(entry_type))
-        return None, None
 
+def entry_from_elt(elt, confkey, entry_type, fields, short_year, use_ee_as_url):
+    """transform a parsed DBLP xml element (article/inproceedings/...) into a dictionnary ready to be output as bibtex"""
     entry = {}
     authors = []  # list of pairs (full author name, last name for BibTeX key)
     pages_error = None
@@ -1023,6 +1071,23 @@ def xml_to_entry(xml, confkey, entry_type, fields, short_year, use_ee_as_url):
     return key, entry
 
 
+def xml_to_entry(xml, confkey, entry_type, fields, short_year, use_ee_as_url):
+    """transform a live-fetched DBLP xml document into a dictionnary ready to be output as bibtex
+
+    (fallback path for keys missing from the local dump, e.g. too recent)
+    """
+    try:
+        tree = XML(xml)
+    except ElementTree.ParseError as e:
+        logging.exception("XML Parsing Error")
+        return None, None
+    elt = tree.find(entry_type.lower())
+    if elt is None:
+        logging.warning('Entry type is not "{0}"'.format(entry_type))
+        return None, None
+    return entry_from_elt(elt, confkey, entry_type, fields, short_year, use_ee_as_url)
+
+
 def write_entry(f, key, entry, entry_type):
     """write the bibtex entry "entry" with key "key" in file "f" """
 
@@ -1173,15 +1238,43 @@ def subs(s):
                 entries[key] = entry
     else:
         # DBLP
+        if not (os.path.exists(DBLP_DUMP_GZ) and os.path.exists(DBLP_DUMP_DTD)):
+            logging.error(
+                'Local DBLP dump not found in "{0}".\n'
+                "Run `python3 fetch_dblp_dump.py` first to download it "
+                "(see db_import/fetch_dblp_dump.py --help).".format(DBLP_DUMP_DIR)
+            )
+            sys.exit(1)
+
+        pubs = []
+        wanted_keys = set()
         for pub in re.finditer(
                 r'href="(https://dblp.uni-trier.de/rec/(?:bibtex/|xml/|)(?:conf|journals)/[^"]*.xml)"',
                 html_conf,
         ):
             url_pub = pub.group(1)
-            logging.info("Parse: <{}>".format(url_pub))
-            xml = get_url(url_pub)
-            key, entry = xml_to_entry(xml, confkey, entry_type, fields_dblp, short_year,
-                                      use_ee_as_url=conf_dict["use_ee_as_url"])
+            m = re.search(r'/rec/(?:bibtex/|xml/|)((?:conf|journals)/[^"]*)\.xml$', url_pub)
+            key_dblp = m.group(1)
+            pubs.append((key_dblp, url_pub))
+            wanted_keys.add(key_dblp)
+
+        logging.info("Looking up {} publication(s) in local DBLP dump...".format(len(wanted_keys)))
+        dtd_entities = load_dtd_entities(DBLP_DUMP_DTD)
+        elements = lookup_dblp_records(DBLP_DUMP_GZ, dtd_entities, wanted_keys)
+
+        for key_dblp, url_pub in pubs:
+            elt = elements.get(key_dblp)
+            if elt is not None:
+                if elt.tag != entry_type.lower():
+                    logging.warning('Entry type of "{0}" is not "{1}"'.format(key_dblp, entry_type))
+                    continue
+                key, entry = entry_from_elt(elt, confkey, entry_type, fields_dblp, short_year,
+                                             use_ee_as_url=conf_dict["use_ee_as_url"])
+            else:
+                logging.info("Not in local dump, falling back to live fetch: <{}>".format(url_pub))
+                xml = get_url(url_pub)
+                key, entry = xml_to_entry(xml, confkey, entry_type, fields_dblp, short_year,
+                                          use_ee_as_url=conf_dict["use_ee_as_url"])
 
             if key is None:
                 continue