tikzpic.py

#!/usr/bin/env python
r"""
    Module to generate a stand-alone PDF of a tikz picture.

    Any lines in the TikZ fragment that begin with "%tikzpic " are preprocessor
    directives to this file.  Presently, only two preprocessor commands are
    supported. They are detailed through the examples below.

    1) To input code into the stand-alone LaTeX before compilation:

        %tikzpic \input{vaucanson.tikz}

    This is useful when multiple figures make use of common TikZ code.

    2) To input verbatim code into the preamble of the stand-alone LaTeX file:

        %tikzpic preamble \usetikzlibrary{positioning}
        %tikzpic preamble \input{cmechabbrev}

    The fragment is placed inside:

        \begin{figure}
            \centering
            % tikz fragment
        \end{figure}

    An example tikz fragment file might look like this:

%tikzpic preamble \input{tikzlibraries}
%tikzpic preamble \input{vaucanson.tikz}
\begin{tikzpicture}[style=vaucanson]
  \node[state] (1) {$\Symbol{1}$};
  \node[state] (0) [right=of 1] {$\Symbol{0}$};
  \path
     (1) edge [bend left]  node {$\frac{1}{2}$} (0)
     (0) edge [bend left]  node {$1$} (1)
     (1) edge [loop left]  node {$\frac{1}{2}$} ();
  \draw (-2.25,0) node {\textbf{(a)}};
\end{tikzpicture}

"""

import os, sys
import re
import shutil
import subprocess
import tempfile

def makedir(dirname):
    import os, errno
    try:
        os.mkdir(dirname)
    except OSError, e:
        if e.errno == errno.EEXIST:
            pass
        else:
            raise


includegraphics_regex = re.compile(r'\\includegraphics(\[.*?\])?\{(.*?)\}')

# Must escape { and }
template = r"""
\documentclass{{article}}

\pagestyle{{empty}}

\usepackage[paper=a0paper]{{geometry}}
\usepackage{{amsmath}}
\usepackage{{amsfonts}}
\usepackage{{graphicx}}
\usepackage{{tikz}}
\usepackage{{comment}}

{preamble}

% If the tikz fragment was used with 'external' tikz library, then it might
% contain some additional commands which are not useful in standalone files.
% So we make the command do nothing.
\def\tikzsetnextfilename#1{{}}

\begin{{document}}

\begin{{figure}}
    \centering
    {code}
\end{{figure}}

\end{{document}}
"""

def default_opener(filename):
    """Opens *filename* using system's default program.

    Parameters
    ----------
    filename : str
        The path of the file to be opened.

    """
    cmds = {'darwin': ['open'],
            'linux2': ['xdg-open'],
            'win32': ['cmd.exe', '/c', 'start', '']}
    cmd = cmds[sys.platform] + [filename]
    subprocess.call(cmd)


def all_images(text, ext='pdf'):
    """Returns a list of all PDFs used by a TeX document.

    The function assumes all tikzpic commands have already been resolved into
    includegraphics commands by create_tex.

    """
    graphics = includegraphics_regex.findall(text)
    basenames = set([ x[1] for x in graphics ])
    images = [x + ".pdf" for x in basenames]
    return images


def create_tex(tikzfile):
    """Returns a string containing the text for a standalone TeX file.

    Parameters
    ----------
    tikzfile : str
        The location of the file containing the TikZ code.

    Returns
    -------
    texcode : str
        The TeX code that can be used to compile a standalone PDF.

    """
    tikzpath = os.path.abspath(tikzfile)
    tikzdir, tikzfile = os.path.split(tikzpath)

    olddir = os.path.abspath(os.path.curdir)
    texcode = ''

    preprocessor = r"%tikzpic "
    query0 = r"\input{"
    query1 = preprocessor + query0
    query2 = preprocessor + "preamble "

    try:
        os.chdir(tikzdir)
        with open(tikzfile, 'r') as fh:
            tikz = []
            preamble = []
            for line in fh.readlines():
                if line.startswith(query1):
                    # Resolve any inputs, usually tikz inputs.
                    tikz.append(line)
                    fname = line.strip()[len(query1):-1]
                    tikz.extend(open(fname, 'r').readlines())
                elif line.startswith(query2):
                    # Handle preamble commands
                    preamble.append(line)
                    cmd = line[len(query2):]
                    if cmd.startswith(query0):
                        # If the command is \input, resolve it.
                        fname = cmd.strip()[len(query0):-1]
                        ext = ".tex"
                        tikzext = ".tikz"
                        if not fname.endswith(ext) and not fname.endswith(tikzext):
                            fname += ext
                        preamble.extend(open(fname, 'r').readlines())
                    else:
                        # Otherwise, put it in verbatim.
                        preamble.append(line[len(query2):])
                else:
                    tikz.append(line)

            texcode = template.format(preamble=''.join(preamble),
                                      code=''.join(tikz))

            # Now go through and do find/replace on includegraphics commands
            def replacer(matchobj):
                options, fn = matchobj.groups()
                fn = os.path.basename(fn)
                if options is None:
                    repl = r"\includegraphics{{{0}}}".format(fn)
                else:
                    repl = r"\includegraphics{0}{{{1}}}".format(options, fn)
                return repl

            texcode = includegraphics_regex.sub(replacer, texcode)

    finally:
        os.chdir(olddir)

    return texcode


def compile_tex(tikzfile, batch=False, prog='lualatex'):
    """Compiles a TikZ figure from a TikZ fragment file.

    Parameters
    ----------
    tikzfile : str
        The location of the file containing the TikZ code.

    Returns
    -------
    pdffile : str
        The path of the PDF of the compiled TikZ figure.

    """
    pdflatex = prog
    pdfcrop = 'pdfcrop'

    # We do not create a new temporary file each time since we want secondary
    # compilations to update an already opened PDF.

    tikzpath = os.path.abspath(tikzfile)
    tikzdir, tikzfn = os.path.split(tikzpath)
    tikzext = tikzfn.split('.')[-1]
    bn = tikzfn[:-len(tikzext)]
    tikzpdf = tikzpath[:-len(tikzext)] + 'pdf'

    prefix = 'tikzpic_'
    suffix = 'tex'
    fname = bn + suffix
    tmpdir = os.path.join(tempfile.gettempdir(), prefix + bn[:-1])
    makedir(tmpdir)
    fpath = os.path.join(tmpdir, fname)
    texdir = os.path.split(fpath)[0]
    tex = create_tex(tikzfile)
    with open(fpath, 'w') as fh:
        fh.write(tex)

    pdffile = fpath[:-3] + 'pdf'
    olddir = os.path.abspath(os.path.curdir)

    # Copy all image dependencies over...
    # This assumes the argument is raw text, not a command.
    try:
        os.chdir(tikzdir)
    except:
        raise
    else:
        with open(tikzfn) as fh:
            images = all_images(fh.read())
        for img in images:
            shutil.copy(img, texdir)
    finally:
        os.chdir(olddir)

    print
    print "Compiling TikZ fragment in {0}".format(fpath)
    print

    try:
        # Keep all intermediate tex files in the same dir as the tex file.
        os.chdir(texdir)
        # Create the pdf.
        if batch:
            args = [pdflatex, '-interaction=batchmode', fname]
        else:
            args = [pdflatex, fname]
        subprocess.call(args)
        subprocess.call(args)
        # Crop it, we must run in a shell to the shebang is interpreted.
        subprocess.call(' '.join([pdfcrop, pdffile, pdffile]), shell=True)
        # Put a copy next to fragment file
        try:
            shutil.copy(pdffile, tikzpdf)
        except IOError as e:
            print e
            tikzpdf = None
    finally:
        os.chdir(olddir)

    return pdffile, tikzpdf

def parse_args(args):
    """Parses the arguments and determines batch and display mode.

    Returns
    -------
    batch : bool
        True if we run in batch mode.
    display : bool
        True if the generated PDFs should be displayed.
    tikzfiles : list
        The list of tikzfiles.

    """
    batch_dict   = {0: False, 1: True, 2: None}
    display_dict = {0: False, 1: True, 2: None}

    batch_arg   = '-b2'
    display_arg = '-d2'
    for i, arg in enumerate( args[1:3] ):
        idx = i+1
        if arg.startswith('-b'):
            batch_arg = arg
            tikzfiles = args[idx+1:]
        elif arg.startswith('-d'):
            display_arg = arg
            tikzfiles = args[idx+1:]
        else:
            # Optional arguments must be specified before files.
            tikzfiles = args[idx:]
            break

    batch = batch_dict.get(int(batch_arg[2]), None)
    display = display_dict.get(int(display_arg[2]), None)

    # Handle defaults if -b and -d are unspecified.
    if len(tikzfiles) < 2:
        if batch is None:
            batch = False
        if display is None:
            display = True
    else:
        if batch is None:
            batch = True
        if display is None:
            display = False

    return batch, display, tikzfiles

if __name__ == '__main__':
    # Should probably use argparse or something similar.
    help = \
"""Usage: python {0} [-b{{n}}] [-d{{n}}] tikzfile.tikz [otherfiles.tikz [...]]

Compiles fragment TikZ files into PDFs.

The optional -b{{n}} parameter controls whether or not LaTeX runs in batch mode.
The {{n}} should be one of the integers 0, 1, or 2.  Default: -b2.

    -b0     Do not run LaTeX with -interaction=batchmode.
    -b1     Run LaTeX with -interaction=batchmode.
    -b2     If a single TikZ file is specified,   -b0 is assumed.
            If multiple TiKZ files are specified, -b1 is assumed.

The optional -d{{n}} parameter controls whether or not the generated PDFs are
displayed. The {{n}} should be one of the integers 0, 1, or 2.  Default: -d2.

    -d0     No generated PDFs are displayed.
    -d1     All generated PDFs are displayed.
    -d2     If multiple TikZ files are specified, -d0 is assumed.
            If a single TikZ file  is  specified, -d1 is assumed.
"""

    if len(sys.argv) < 2:
        print help.format(sys.argv[0])
    else:
        batch, display, tikzfiles = parse_args(sys.argv)
        if len(tikzfiles) == 0:
            print help.format(sys.argv[0])
        else:
            for tikzfile in tikzfiles:
                pdffile, tikzpdf = compile_tex(tikzfile, batch=batch)
                print "\nTikZ fragment:\n\t{0}".format(tikzfile)
                print "\nCompiled TeX file:\n\t{0}".format(pdffile[:-3] + 'tex')
                print "\nOutput PDF file:\n\t{0}".format(pdffile)
                if tikzpdf is not None:
                    print "\t{0}".format(tikzpdf)
                print
                if display and pdffile is not None:
                    default_opener(pdffile)