Source code for schrodinger.application.desmond.xvfbwrapper

#!/usr/bin/env python
#
#   * Corey Goldberg, 2012, 2013
#
#   * inspired by:
#       PyVirtualDisplay: http://pypi.python.org/pypi/PyVirtualDisplay
"""wrapper for running display inside X virtual framebuffer (Xvfb)"""

import fnmatch
import os
import random
import subprocess
import time


class Xvfb:

    def __init__(self, width=800, height=680, colordepth=24, **kwargs):
        self.width = width
        self.height = height
        self.colordepth = colordepth
        self._set_xvfb_cmd()

        for key, value in kwargs.items():
            self.xvfb_cmd = self.xvfb_cmd + ['-%s' % key, value]

        self.proc = None
        if 'DISPLAY' in os.environ:
            self.old_display_num = os.environ['DISPLAY'].split(':')[1]
        else:
            self.old_display_num = None

    def _set_xvfb_cmd(self):
        self.xvfb_cmd = [
            '-screen', '0',
            '%dx%dx%d' % (self.width, self.height, self.colordepth)
        ]

    def __enter__(self):
        self.start()
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        self.stop()

    def start(self):
        if 'DISPLAY' in os.environ:
            print("DISPLAY already exists: %s " % os.environ['DISPLAY'])
            return
        self.vdisplay_num = self.search_for_free_display()
        self.xvfb_cmd = ['Xvfb', ':%d' % self.vdisplay_num] + self.xvfb_cmd

        self.proc = subprocess.Popen(self.xvfb_cmd,
                                     stdout=subprocess.DEVNULL,
                                     stderr=subprocess.DEVNULL,
                                     universal_newlines=True)
        print(self.xvfb_cmd)
        time.sleep(2)  # give Xvfb time to start
        ret_code = self.proc.poll()
        if ret_code is None:
            self._redirect_display(self.vdisplay_num)
        elif self.old_display_num:
            self._redirect_display(self.old_display_num)
            self.proc = None
            print('Error: Xvfb did not start')

    def stop(self):
        if self.old_display_num:
            self._redirect_display(self.old_display_num)
        elif 'DISPLAY' in os.environ:
            del os.environ['DISPLAY']
        if self.proc is not None:
            self.proc.kill()
            self.proc.wait()
            self.proc = None
        self._set_xvfb_cmd()

    def __del__(self):
        print("deleting X...")
        self.stop()
        del self

    def search_for_free_display(self):
        ls = [int(x.split('X')[1].split('-')[0]) for x in self._lock_files()]
        min_display_num = 1000
        if len(ls):
            display_num = max(min_display_num, max(ls) + 1)
        else:
            display_num = min_display_num
        random.seed()
        display_num += random.randint(0, 100)
        return display_num

    def _lock_files(self):
        tmpdir = '/tmp'
        pattern = '.X*-lock'
        names = fnmatch.filter(os.listdir(tmpdir), pattern)
        ls = [os.path.join(tmpdir, child) for child in names]
        ls = [p for p in ls if os.path.isfile(p)]
        return ls

    def _redirect_display(self, display_num):
        if display_num:
            os.environ['DISPLAY'] = ':%s' % display_num