Use luatex's kpsewhich for speed.

anntzer · anntzer · commit 5e13e0e49d0d · 2021-02-20T23:04:26.000+01:00
The machinery for running an interactive tex session mostly already
existed in backend_pgf, so directly reuse that (as a base class).  On
the matplotlib macos, this significantly speeds up
```sh
python -c 'from pylab import *; mpl.use("pdf"); rcParams["text.usetex"] = True; plot(); savefig("/tmp/test.pdf", backend="pdf")'
```
from ~4.5s to ~2.5s.

Note that filesystem encodings may be a bit iffy, we still need to check
how things go through the various layers here; it may end up being best
to make the process streams be binary and perform the encoding/decoding
ourselves.

We also need to figure out how to best advertise this (do we emit a
warning suggesting to install luatex on windows and macos if luatex is
not present?).
diff --git a/lib/matplotlib/backends/backend_pgf.py b/lib/matplotlib/backends/backend_pgf.py
@@ -16,7 +16,7 @@
 from PIL import Image
 
 import matplotlib as mpl
-from matplotlib import _api, cbook, font_manager as fm
+from matplotlib import _api, cbook, font_manager as fm, texmanager
 from matplotlib.backend_bases import (
     _Backend, _check_savefig_extra_args, FigureCanvasBase, FigureManagerBase,
     GraphicsContextBase, RendererBase, _no_output_draw
@@ -188,25 +188,23 @@ def gs_convert(pdffile, pngfile, dpi):
     raise RuntimeError("No suitable pdf to png renderer found.")
 
 
-class LatexError(Exception):
-    def __init__(self, message, latex_output=""):
-        super().__init__(message)
-        self.latex_output = latex_output
+LatexError = texmanager._InteractiveTex.TexError
 
-    def __str__(self):
-        s, = self.args
-        if self.latex_output:
-            s += "\n" + self.latex_output
-        return s
 
-
-class LatexManager:
+class LatexManager(texmanager._InteractiveTex):
     """
     The LatexManager opens an instance of the LaTeX application for
     determining the metrics of text elements. The LaTeX environment can be
     modified by setting fonts and/or a custom preamble in `.rcParams`.
     """
 
+    tmpdir = property(lambda self: self._tmpdir.name)
+    texcommand = property(lambda self: self._texcmd)
+    latex = property(
+        lambda self: self._tex if self._tex.returncode is not None else None)
+    latex_stdin_utf8 = _api.deprecated("3.3")(
+        property(lambda self: self._tex.stdin))
+
     @staticmethod
     def _build_latex_header():
         latex_preamble = get_preamble()
@@ -225,7 +223,6 @@ def _build_latex_header():
             latex_fontspec,
             r"\begin{document}",
             r"text $math \mu$",  # force latex to load fonts now
-            r"\typeout{pgf_backend_query_start}",
         ]
         return "\n".join(latex_header)
 
@@ -242,87 +239,10 @@ def _get_cached_or_new(cls):
     def _get_cached_or_new_impl(cls, header):  # Helper for _get_cached_or_new.
         return cls()
 
-    def _stdin_writeln(self, s):
-        if self.latex is None:
-            self._setup_latex_process()
-        self.latex.stdin.write(s)
-        self.latex.stdin.write("\n")
-        self.latex.stdin.flush()
-
-    def _expect(self, s):
-        s = list(s)
-        chars = []
-        while True:
-            c = self.latex.stdout.read(1)
-            chars.append(c)
-            if chars[-len(s):] == s:
-                break
-            if not c:
-                self.latex.kill()
-                self.latex = None
-                raise LatexError("LaTeX process halted", "".join(chars))
-        return "".join(chars)
-
-    def _expect_prompt(self):
-        return self._expect("\n*")
-
     def __init__(self):
-        # create a tmp directory for running latex, register it for deletion
-        self._tmpdir = TemporaryDirectory()
-        self.tmpdir = self._tmpdir.name
-        self._finalize_tmpdir = weakref.finalize(self, self._tmpdir.cleanup)
-
-        # test the LaTeX setup to ensure a clean startup of the subprocess
-        self.texcommand = mpl.rcParams["pgf.texsystem"]
-        self.latex_header = LatexManager._build_latex_header()
-        latex_end = "\n\\makeatletter\n\\@@end\n"
-        try:
-            latex = subprocess.Popen(
-                [self.texcommand, "-halt-on-error"],
-                stdin=subprocess.PIPE, stdout=subprocess.PIPE,
-                encoding="utf-8", cwd=self.tmpdir)
-        except FileNotFoundError as err:
-            raise RuntimeError(
-                f"{self.texcommand} not found.  Install it or change "
-                f"rcParams['pgf.texsystem'] to an available TeX "
-                f"implementation.") from err
-        except OSError as err:
-            raise RuntimeError("Error starting process %r" %
-                               self.texcommand) from err
-        test_input = self.latex_header + latex_end
-        stdout, stderr = latex.communicate(test_input)
-        if latex.returncode != 0:
-            raise LatexError("LaTeX returned an error, probably missing font "
-                             "or error in preamble.", stdout)
-
-        self.latex = None  # Will be set up on first use.
         self.str_cache = {}  # cache for strings already processed
-
-    def _setup_latex_process(self):
-        # Open LaTeX process for real work; register it for deletion.  On
-        # Windows, we must ensure that the subprocess has quit before being
-        # able to delete the tmpdir in which it runs; in order to do so, we
-        # must first `kill()` it, and then `communicate()` with it.
-        self.latex = subprocess.Popen(
-            [self.texcommand, "-halt-on-error"],
-            stdin=subprocess.PIPE, stdout=subprocess.PIPE,
-            encoding="utf-8", cwd=self.tmpdir)
-
-        def finalize_latex(latex):
-            latex.kill()
-            latex.communicate()
-
-        self._finalize_latex = weakref.finalize(
-            self, finalize_latex, self.latex)
-        # write header with 'pgf_backend_query_start' token
-        self._stdin_writeln(self._build_latex_header())
-        # read all lines until our 'pgf_backend_query_start' token appears
-        self._expect("*pgf_backend_query_start")
-        self._expect_prompt()
-
-    @_api.deprecated("3.3")
-    def latex_stdin_utf8(self):
-        return self.latex.stdin
+        self.latex_header = LatexManager._build_latex_header()
+        super().__init__(mpl.rcParams["pgf.texsystem"], self.latex_header)
 
     def get_width_height_descent(self, text, prop):
         """
@@ -347,7 +267,7 @@ def get_width_height_descent(self, text, prop):
                              .format(text, e.latex_output)) from e
 
         # typeout width, height and text offset of the last textbox
-        self._stdin_writeln(r"\typeout{\the\wd0,\the\ht0,\the\dp0}")
+        self._stdin_writeln(r"\message{\the\wd0,\the\ht0,\the\dp0}")
         # read answer from latex and advance to the next prompt
         try:
             answer = self._expect_prompt()
@@ -357,7 +277,7 @@ def get_width_height_descent(self, text, prop):
 
         # parse metrics from the answer string
         try:
-            width, height, offset = answer.splitlines()[0].split(",")
+            width, height, offset = answer.split(",")
         except Exception as err:
             raise ValueError("Error processing '{}'\nLaTeX Output:\n{}"
                              .format(text, answer)) from err
diff --git a/lib/matplotlib/dviread.py b/lib/matplotlib/dviread.py
@@ -30,7 +30,7 @@
 
 import numpy as np
 
-from matplotlib import _api, cbook, rcParams
+from matplotlib import _api, cbook, texmanager, rcParams
 
 _log = logging.getLogger(__name__)
 
@@ -1040,6 +1040,30 @@ def _parse_enc(path):
             "Failed to parse {} as Postscript encoding".format(path))
 
 
+class _LuatexKpathsea(texmanager._InteractiveTex):
+    @lru_cache
+    def __new__(cls):
+        self = super().__new__(cls)
+        super(cls, self).__init__("luatex")
+        return self
+
+    def __init__(self):
+        pass  # Skip the super().__init__.
+
+    def find_tex_file(self, filename, format=None):
+        if format is not None:
+            filename = f"{filename}.{format}"
+        if "\\" in filename or "'" in filename:  # Skip unrealistic escapes.
+            raise ValueError(f"Invalid filename: {filename!r}")
+        # While kpse.find_file seems appropriate, it doesn't actually handle
+        # extensions in the filename (contrary to its docs), and mapping
+        # extension to format names (".pfb" -> "type1 fonts") would require a
+        # large hard-coded table.
+        self._stdin_writeln(r"\directlua{print(kpse.lookup('%s'))}" % filename)
+        val = self._expect_prompt().rstrip()
+        return "" if val == "nil" else val
+
+
 @lru_cache()
 def find_tex_file(filename, format=None):
     """
@@ -1072,6 +1096,13 @@ def find_tex_file(filename, format=None):
     if isinstance(format, bytes):
         format = format.decode('utf-8', errors='replace')
 
+    try:
+        lk = _LuatexKpathsea()
+    except FileNotFoundError:
+        pass
+    else:
+        return lk.find_tex_file(filename, format)
+
     if os.name == 'nt':
         # On Windows only, kpathsea can use utf-8 for cmd args and output.
         # The `command_line_encoding` environment variable is set to force it
diff --git a/lib/matplotlib/texmanager.py b/lib/matplotlib/texmanager.py
@@ -29,6 +29,7 @@
 import re
 import subprocess
 from tempfile import TemporaryDirectory
+import weakref
 
 import numpy as np
 
@@ -401,3 +402,73 @@ def get_text_width_height_descent(self, tex, fontsize, renderer=None):
                 page, = dvi
             # A total height (including the descent) needs to be returned.
             return page.width, page.height + page.descent, page.descent
+
+
+class _InteractiveTex:
+    """
+    Interactive tex process supporting communication through stdin and stdout.
+    """
+
+    class TexError(Exception):
+        def __init__(self, message, latex_output=""):
+            super().__init__(message)
+            self.latex_output = latex_output
+
+        def __str__(self):
+            s, = self.args
+            if self.latex_output:
+                s += "\n" + self.latex_output
+            return s
+
+    def __init__(self, cmd, header=""):
+        self._tmpdir = TemporaryDirectory()
+        self._finalize_tmpdir = weakref.finalize(self, self._tmpdir.cleanup)
+        self._texcmd = cmd
+        self._header = header
+        self._setup_tex_process(cmd, header)
+
+    def _stdin_writeln(self, s):
+        if self._tex.returncode is not None:
+            self._setup_tex_process(self._texcmd, self._header)
+        self._tex.stdin.write(s)
+        self._tex.stdin.write("\n")
+        self._tex.stdin.flush()
+
+    def _expect(self, s):
+        s = list(s)
+        chars = []
+        while True:
+            c = self._tex.stdout.read(1)
+            chars.append(c)
+            if chars[-len(s):] == s:
+                break
+            if not c:
+                self._tex.kill()
+                raise self.TexError("TeX process halted", "".join(chars))
+        return "".join(chars)
+
+    def _expect_prompt(self):
+        return self._expect("\n*")[:-2]
+
+    def _setup_tex_process(self, cmd, header):
+        # Open TeX process; register it for deletion.  On Windows, we must
+        # ensure that the subprocess has quit before being able to delete the
+        # tmpdir in which it runs; in order to do so, we must first `kill()`
+        # it, and then `communicate()` with it.
+        self._tex = subprocess.Popen(
+            [cmd, "-halt-on-error"],
+            stdin=subprocess.PIPE, stdout=subprocess.PIPE,
+            encoding="utf-8", cwd=self._tmpdir.name)
+
+        def finalize_tex(tex):
+            tex.kill()
+            tex.communicate()
+
+        self._finalize_tex = weakref.finalize(self, finalize_tex, self._tex)
+
+        self._stdin_writeln(header)
+        # Emit a marker once the header is handled, and wait for it to appear.
+        marker = "init-done"
+        self._stdin_writeln(r"\immediate\write17{%s}" % marker)
+        self._expect("*%s" % marker)
+        self._expect_prompt()
