Skip to content
Snippets Groups Projects
pipeline.py 4.92 KiB
Newer Older
#!/usr/bin/env python3

James Vasile's avatar
James Vasile committed
import sys
import re
import click
import frontmatter
import importlib
import os
import pkgutil
import pprint
pp = pprint.PrettyPrinter(indent=4)
"""
The 'stage' option is for running this pipeline either pre- or post- latex.
Ex: pipeline.py target.ltx -o stage post
Ex: pipeline.py target.ltx -o stage pre
"""
James Vasile's avatar
James Vasile committed

def dump(text, meta, fname="/tmp/t.tex"):
    """Writes the current text and meta to /tmp/t.tex or fname if
       specified"""
James Vasile's avatar
James Vasile committed
    with open(fname, 'w') as fh:
        fh.write(pp.pformat(meta))
        fh.write("\n---\n")
        fh.write(text)
James Vasile's avatar
James Vasile committed


def err(msg):
    """Yeah, this is an obnoxious error message, but the user will have to
    pick it out of a long scroll of LaTeX and make output.

    """
    sys.stderr.write("BEGIN PIPELINE ERROR MSG\n")
    sys.stderr.write(msg)
    sys.stderr.write("\n")
    sys.stderr.write("END PIPELINE ERROR MSG\n")
    sys.exit(-1)
James Vasile's avatar
James Vasile committed


def get_plugins(plugin_dir=None):
    """Load plugins from PLUGIN_DIR and return a dict with the plugin name
    hashed to the imported plugin.

    PLUGIN_DIR is the name of the dir from which to load plugins.  If
    it is None, use the plugin dir in the dir that holds this func.

    We load plugins and run them in asciibetical order.  We ignore
    plugins that begin with a character other than a digit or a
    letter.
    run_p(text, meta, opt): predicate returns True if this
    plugin thinks it should run in the pipeline.

    run(text, meta, opt): runs the plugin, returns text, meta

    after_p(pdf_fname, meta, opt): predicate returns True if this plugin
    thinks it should run after the pdf is produced.

    after(pdf_fname, meta, opt): runs the plugin, returns meta.  May change the pdf

    """
    if not plugin_dir:
        plugin_dir = os.path.join(
            os.path.dirname(
                os.path.abspath(__file__)), "plugins")

    plugins = {}
    pat = {
        "enabled": r"^[0-9a-zA-Z]",
        "prefix": r"^[0-9]+_",
    }
James Vasile's avatar
James Vasile committed
    pat = {k: re.compile(v) for k, v in pat.items()}
    for fname in sorted(os.listdir(plugin_dir)):
        if (not fname.endswith(".py")
            or fname.startswith('.')
            or '#' in fname
James Vasile's avatar
James Vasile committed
                or not pat['enabled'].match(fname)):
James Vasile's avatar
James Vasile committed
        spec = importlib.util.spec_from_file_location(
            fname, os.path.join(plugin_dir, fname))
        fname = pat["prefix"].sub("", fname)
        plugins[fname] = importlib.util.module_from_spec(spec)
        spec.loader.exec_module(plugins[fname])
    return plugins

James Vasile's avatar
James Vasile committed

def tuple2dict(tup):
James Vasile's avatar
James Vasile committed
    """TUP is an iterable of 2-tuples.  Return a dict where the first element
    of each tuple hashes to the second.

    """
    ret = {}
    for t in tup:
        ret[t[0]] = t[1]
    return ret
James Vasile's avatar
James Vasile committed


class OTSFrontmatterError(Exception):
    """Error reading a document's frontmatter."""
    def __init__(self, doc_name):
        self.doc_name = doc_name
    def __str__(self):
        return f"""ERROR: problem reading frontmatter

        In '{self.doc_name}':

        This error usually comes from an unterminated string constant,
        e.g., something like:

        ---
        title: "Effects of Forgetting a Closing Double-Quote
        date: 24 July 2023
        draft: false
        ---\n"""


@click.command()
@click.argument('filename')
@click.option('--output', help="output filename")
James Vasile's avatar
James Vasile committed
@click.option(
    '--plugin',
    help="specify just one plugin to run (not implemented)")
@click.option(
    '--option',
    '-o',
    nargs=2,
    multiple=True,
    default=None,
James Vasile's avatar
James Vasile committed
    help='Variables and their values.')
def cli(filename, output, option, plugin):
James Vasile's avatar
James Vasile committed
    pdf_fname = os.path.splitext(filename)[0] + '.pdf'
    try:
        doc = frontmatter.load(filename)
    except Exception:
        raise OTSFrontmatterError(filename)
    text = doc.content
    meta = doc.metadata
James Vasile's avatar
James Vasile committed
    if meta == {}:
        meta['legacy'] = "legacy"
    meta['input_filename'] = filename
    meta['output_filename'] = output
James Vasile's avatar
James Vasile committed

    # Grab the environment too
James Vasile's avatar
James Vasile committed
    meta['environment'] = dict(os.environ)

    # Pass to jinja any options submitted on the command line
    for o in option:
        meta[o[0]] = o[1]

    # Dict is easier to work with
    option = tuple2dict(option)
    option['stage'] = option.get('stage', 'pre')
James Vasile's avatar
James Vasile committed

    plugins = get_plugins()
James Vasile's avatar
James Vasile committed
    for p, m in plugins.items():
        if option['stage'] == 'pre':
            if hasattr(m, "run_p"):
                if m.run_p(text, meta):
                    text, meta = m.run(text, meta)
        elif option['stage'] == 'post':
            if hasattr(m, "after_p"):
                if m.after_p(pdf_fname, meta):
                    meta = m.after(pdf_fname, meta)
        else:
            err("Unknown stage: %s" % option['stage'])

James Vasile's avatar
James Vasile committed
    # Print the output or write it to a file if we're pre-latex
    if option['stage'] == 'pre':
        if output:
            with open(output, 'w') as fh:
                fh.write(text)
        else:
            print(text)
James Vasile's avatar
James Vasile committed


if __name__ == '__main__':
    cli()