Improve headlessness detection for backend selection.

anntzer · anntzer · commit 1a88ee6c75a8 · 2020-05-13T15:29:14.000+02:00
We currently check the $DISPLAY environment variable to autodetect
whether we should auto-pick a non-interactive backend on Linux, but that
variable can be set to an "invalid" value.  A realistic use case is for
example a tmux session started interactively inheriting an initially
valid $DISPLAY, but to which one later reconnects e.g. via ssh, at which
point $DISPLAY becomes invalid.

Before this PR, something like
```
DISPLAY=:123 MPLBACKEND= MATPLOTLIBRC=/dev/null python -c 'import pylab'
```
(where we unset matplotlibrc to force backend autoselection) would crash
when we select qt and qt fails to initialize as $DISPLAY is invalid (qt
unconditionally abort()s via qFatal() in that case).
With this PR, we correctly autoselect a non-interactive backend.

Note that we were *already* relying on $DISPLAY being correctly set
before, and this PR also doesn't return "invalid display" if we can't
load X11, so this should not make anything worse on Wayland (at worst
we'll just fail to detect headlessness like before).
diff --git a/lib/matplotlib/backends/backend_qt5.py b/lib/matplotlib/backends/backend_qt5.py
@@ -7,7 +7,6 @@
 import traceback
 
 import matplotlib
-
 from matplotlib import backend_tools, cbook
 from matplotlib._pylab_helpers import Gcf
 from matplotlib.backend_bases import (
@@ -115,10 +114,8 @@ def _create_qApp():
                     is_x11_build = False
             else:
                 is_x11_build = hasattr(QtGui, "QX11Info")
-            if is_x11_build:
-                display = os.environ.get('DISPLAY')
-                if display is None or not re.search(r':\d', display):
-                    raise RuntimeError('Invalid DISPLAY variable')
+            if is_x11_build and matplotlib._c_internal_utils.invalid_display():
+                raise RuntimeError('Invalid DISPLAY variable')
 
             try:
                 QtWidgets.QApplication.setAttribute(
diff --git a/lib/matplotlib/cbook/__init__.py b/lib/matplotlib/cbook/__init__.py
@@ -28,6 +28,8 @@
 import numpy as np
 
 import matplotlib
+import matplotlib as mpl
+import matplotlib._c_internal_utils
 from .deprecation import (
     deprecated, warn_deprecated,
     _rename_parameter, _delete_parameter, _make_keyword_only,
@@ -71,7 +73,8 @@ def _get_running_interactive_framework():
     if 'matplotlib.backends._macosx' in sys.modules:
         if sys.modules["matplotlib.backends._macosx"].event_loop_is_running():
             return "macosx"
-    if sys.platform.startswith("linux") and not os.environ.get("DISPLAY"):
+    if (sys.platform.startswith("linux")
+            and mpl._c_internal_utils.invalid_display()):
         return "headless"
     return None
 
diff --git a/setupext.py b/setupext.py
@@ -422,6 +422,11 @@ def get_extensions(self):
             include_dirs=["extern"])
         add_numpy_flags(ext)
         yield ext
+        # c_internal_utils
+        ext = Extension(
+            "matplotlib._c_internal_utils", ["src/_c_internal_utils.cpp"],
+            libraries={"linux": ["dl"]}.get(sys.platform, []))
+        yield ext
 
 
 class SampleData(OptionalPackage):
diff --git a/src/_c_internal_utils.cpp b/src/_c_internal_utils.cpp
@@ -0,0 +1,55 @@
+#define PY_SSIZE_T_CLEAN
+#include <Python.h>
+
+#ifdef __linux__
+#include <dlfcn.h>
+struct Display;
+typedef Display* (* XOpenDisplay_t)(char* display_name);
+typedef int (* XCloseDisplay_t)(Display* display);
+#endif
+
+static PyObject* mpl_invalid_display(PyObject* /* module */)
+{
+#ifdef __linux__
+  void* x11;
+  XOpenDisplay_t XOpenDisplay;
+  XCloseDisplay_t XCloseDisplay;
+  if (!(x11 = dlopen("libX11.so", RTLD_LAZY)) ||
+      !(XOpenDisplay = (XOpenDisplay_t)dlsym(x11, "XOpenDisplay")) ||
+      !(XCloseDisplay = (XCloseDisplay_t)dlsym(x11, "XCloseDisplay"))) {
+    Py_RETURN_FALSE;
+  }
+  Display* display = XOpenDisplay(NULL);
+  if (display) {
+    XCloseDisplay(display);
+  }
+  if (dlclose(x11)) {
+    PyErr_SetString(PyExc_RuntimeError, dlerror());
+    return NULL;
+  }
+  if (display) {
+    Py_RETURN_FALSE;
+  } else {
+    Py_RETURN_TRUE;
+  }
+#else
+  Py_RETURN_FALSE;
+#endif
+}
+
+static char const* mpl_invalid_display__doc__ =
+"invalid_display()\n--\n\n"
+"Attempt to check whether the current X11 display is invalid.\n\n"
+"Returns True if running on Linux and libX11 can be loaded and \n"
+"XOpenDisplay(NULL) returns NULL, False otherwise.";
+
+static PyMethodDef functions[] = {
+  {"invalid_display", (PyCFunction)mpl_invalid_display, METH_NOARGS, mpl_invalid_display__doc__},
+  {NULL, NULL}};  // sentinel.
+static PyModuleDef util_module = {
+  PyModuleDef_HEAD_INIT, "_c_internal_utils", "", 0, functions, NULL, NULL, NULL, NULL};
+
+PyMODINIT_FUNC PyInit__c_internal_utils(void)
+{
+  return PyModule_Create(&util_module);
+}
