bibx

#!/usr/bin/env python3

import argparse
import bibtexparser
import subprocess
import shlex
import os
from bibtexparser.bparser import BibTexParser, BibDatabase
from bibtexparser.bwriter import BibTexWriter

VERSION = "v1.1"


def get_parser():
    parser = argparse.ArgumentParser(
        description="""Extract bibliography entries from LaTeX sources
            using a reference BibTeX file""",
        epilog="""For reporting bugs, giving suggestions and contributing
            to this project, visit https://github.com/CodePurble/bibextract""",
    )
    parser.add_argument(
        "-v", "--version", help="Show version information and exit", action="store_true"
    )
    parser.add_argument(
        "-o",
        "--output",
        default="ext.bib",
        help="""File to which the extracted entries are output
                        to. Default is './ext.bib'""",
    )
    parser.add_argument(
        "-e",
        "--entries",
        help="""Text file containing entries to extract from
                        the BibTeX file, each line in the file must have ONE
                        BibTeX label""",
    )
    parser.add_argument(
        "-q",
        "--quiet",
        default=False,
        help="Suppress output to stdout",
        action="store_true",
    )
    parser.add_argument(
        "-f",
        "--files",
        help="""LaTeX files to scan for citations whose BibTeX
                        entries need to be extracted""",
        action="extend",
        nargs="+",
    )
    parser.add_argument(
        "-i",
        "--indent",
        help="""Indentation to be used in the generated BibTeX
                        file. Default is four spaces""",
        default="    ",
    )
    parser.add_argument(
        "-d",
        "--dryrun",
        help="""Perform a dry-run, i.e. do not output a file,
                        just print info messages. Ignores the -o option,
                        affected by -q""",
        action="store_true",
    )
    parser.add_argument(
        "-b",
        "--bib",
        help="BibTeX file(s) to look for entries in",
        action="extend",
        nargs="+",
    )
    return parser


def main():
    parser = get_parser()
    args = parser.parse_args()

    if args.version:
        parser.exit(0, f"{VERSION}\n")
    else:
        if (args.bib is None) or (args.files is None):
            parser.print_usage()
            parser.exit(
                2, "Error: the following arguments are required: -b/--bib, -f/--files\n"
            )
        else:
            # Scan LaTeX files for citations and generate set containing BibTeX entry
            # labels
            e_set = set()
            for file in args.files:
                output = subprocess.run(
                    [
                        os.path.dirname(os.path.realpath(__file__)) + "/ext.sh",
                        shlex.quote(file),
                    ],
                    capture_output=True,
                )
                if output.returncode == 0:
                    e_set = e_set.union(
                        set(output.stdout.decode().rstrip().split("\n"))
                    )

            # Merge custom BibTeX entry labels into main set if any
            if args.entries is not None:
                with open(args.entries, "r") as en:
                    while entry := en.readline():
                        e_set = e_set.union(entry)

            e_dict = {}
            for file in args.bib:
                with open(file) as bibtex_file:
                    parser = BibTexParser()
                    parser.ignore_nonstandard_types = False
                    bib_database = bibtexparser.load(bibtex_file, parser)
                    curr_dict = bib_database.entries_dict
                    e_dict.update(curr_dict)

            out_db = BibDatabase()
            found = 0
            for entry in e_set:
                if entry in e_dict.keys():
                    if not args.quiet:
                        print(f"Found: {entry}")
                    found += 1
                    out_db.entries.append(e_dict[entry])
                else:
                    if not args.quiet:
                        print(f"Not found: {entry}")

            writer = BibTexWriter()
            writer.indent = args.indent
            writer.add_trailing_comma = True
            if not args.dryrun:
                with open(args.output, "w") as outfile:
                    bibtexparser.dump(out_db, outfile, writer)

            if not args.quiet:
                print(f"\nFound {found}/{len(e_set)} entries")


if __name__ == "__main__":
    main()