nvidia-prime 0.6.2 / usr / bin / prime-select

#!/usr/bin/python
#
#       prime-select
#
#       Copyright 2013 Canonical Ltd.
#       Author: Alberto Milone <alberto.milone@canonical.com>
#
#       Script to switch between NVIDIA and Intel graphics driver libraries.
#
#       Usage:
#           prime-select   nvidia|intel|query
#           nvidia:   switches to NVIDIA's version of libGL.so
#           intel: switches to the open-source version of libGL.so
#           query: checks which version is currently active and writes
#                  "nvidia", "intel" or "unknown" to the standard output
#
#       Permission is hereby granted, free of charge, to any person
#       obtaining a copy of this software and associated documentation
#       files (the "Software"), to deal in the Software without
#       restriction, including without limitation the rights to use,
#       copy, modify, merge, publish, distribute, sublicense, and/or sell
#       copies of the Software, and to permit persons to whom the
#       Software is furnished to do so, subject to the following
#       conditions:
#
#       The above copyright notice and this permission notice shall be
#       included in all copies or substantial portions of the Software.
#
#       THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND,
#       EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES
#       OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND
#       NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT
#       HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY,
#       WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
#       FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR
#       OTHER DEALINGS IN THE SOFTWARE.


import os
import sys
import re
import subprocess
from subprocess import Popen, PIPE, CalledProcessError

class Alternatives:

    def __init__(self, master_link):
        self._open_drivers_alternative = 'mesa/ld.so.conf'
        self._open_egl_drivers_alternative = 'mesa-egl/ld.so.conf'
        self._command = '/usr/bin/update-alternatives'
        self._master_link = master_link

        # Make sure that the PATH environment variable is set
        if not os.environ.get('PATH'):
            os.environ['PATH'] = '/sbin:/usr/sbin:/bin:/usr/bin'

    def list_alternatives(self):
        '''Get the list of alternatives for the master link'''
        dev_null = open('/dev/null', 'w')
        alternatives = []
        p1 = Popen([self._command, '--list', self._master_link],
                   stdout=PIPE, stderr=dev_null)
        p = p1.communicate()[0]
        dev_null.close()
        c = p.split('\n')
        for line in c:
            line.strip() and alternatives.append(line.strip())

        return alternatives

    def get_current_alternative(self, check_status=False):
        '''Get the alternative in use

        If check_status is True, then only the status will
        be reported (i.e. either "manual" or "auto")
        '''
        dev_null = open('/dev/null', 'w')
        current_alternative = None
        p1 = Popen([self._command, '--query', self._master_link],
                   stdout=PIPE, stderr=dev_null)
        p = p1.communicate()[0]
        dev_null.close()
        c = p.split('\n')
        if check_status:
            pattern = 'Status:'
        else:
            pattern = 'Value:'
        for line in c:
            if line.strip().startswith(pattern):
                return line.replace(pattern, '').strip()
        return None

    def get_alternative_by_name(self, name, ignore_pattern=None):
        '''Get the alternative link by providing the driver name

        ignore_pattern allows ignoring a substring in the name'''
        if ignore_pattern:
            name = name.replace(ignore_pattern, '')
        alternatives = self.list_alternatives()

        for alternative in alternatives:
            if alternative.split('/')[-2] == name:
                return alternative

        return None

    def get_open_drivers_alternative(self):
        '''Get the alternative link for open drivers'''
        return self.get_alternative_by_name(self._open_drivers_alternative)

    def get_open_egl_drivers_alternative(self):
        '''Get the alternative link for open EGL/GLES drivers'''
        return self.get_alternative_by_name(self._open_egl_drivers_alternative)

    def set_alternative(self, path):
        '''Tries to set an alternative and returns the boolean exit status'''
        try:
            subprocess.check_call([self._command, '--set',
                                   self._master_link, path])
            self.ldconfig()
        except CalledProcessError:
            return False
        return True

    def ldconfig(self):
        '''Call ldconfig'''
        try:
            subprocess.check_call(['/sbin/ldconfig'])
        except CalledProcessError:
            return False
        return True


class Switcher(object):

    def __init__(self):
        self._power_profile_path = '/etc/prime-discrete'
        self._supported_architectures = {'i386': 'i386', 'amd64': 'x86_64'}
        main_arch = self._get_architecture()
        other_arch = self._supported_architectures.values()[
                          int(not self._supported_architectures
                          .values().index(main_arch))]
        main_alternative_name = self._get_alternative_name_from_arch(main_arch)
        other_alternative_name = self._get_alternative_name_from_arch(other_arch)

        # We have 2 alternatives, one for each architecture
        self._gl_switcher = Alternatives(main_alternative_name)
        self._gl_switcher_other = Alternatives(other_alternative_name)

    def _get_architecture(self):
        dev_null = open('/dev/null', 'w')
        p1 = Popen(['dpkg', '--print-architecture'], stdout=PIPE, stderr=dev_null)
        p = p1.communicate()[0]
        dev_null.close()
        architecture = p.strip()
        return self._supported_architectures.get(architecture)

    def _get_alternative_name_from_arch(self, architecture):
        alternative = '%s-linux-gnu_gl_conf' % architecture
        return alternative

    def _simplify_x_alternative_name(self, alternative):
        return alternative.split('/')[-1]

    def _simplify_gl_alternative_name(self, alternative):
        return alternative.split('/')[-2]

    def _get_gl_alternatives_list(self):
        raw_alternatives = self._gl_switcher.list_alternatives()
        return  [ self._simplify_gl_alternative_name(x) for x in raw_alternatives ]

    def _update_initramfs(self):
        subprocess.call(['update-initramfs', '-u'])
        # This may not be necessary
        subprocess.call(['update-initramfs', '-u', '-k', os.uname()[2]])

    def _get_current_alternative(self):
        # This is a list as the 2nd item belongs to another architecture
        alternatives = [None, None]
        raw_gl_alternative = self._gl_switcher.get_current_alternative()
        raw_gl_alternative_other = self._gl_switcher_other.get_current_alternative()

        if (raw_gl_alternative != None):
            alternatives[0] = self._simplify_gl_alternative_name(raw_gl_alternative)
        if (raw_gl_alternative_other != None):
            alternatives[1] = self._simplify_gl_alternative_name(raw_gl_alternative_other)

        return alternatives

    def _get_current_alternative_status(self):
        '''Get the current status i.e. either "auto" or "manual"'''
        # This is a list as the 2nd item belongs to another architecture
        alternatives = [None, None]
        raw_gl_alternative = self._gl_switcher.get_current_alternative(True)
        raw_gl_alternative_other = self._gl_switcher_other.get_current_alternative(True)

        if (raw_gl_alternative != None):
            alternatives[0] = raw_gl_alternative
        if (raw_gl_alternative_other != None):
            alternatives[1] = raw_gl_alternative_other

        return alternatives

    def enable_alternative(self, alternative):
        # Make sure that the alternative exists
        gl_alternative = self._gl_switcher.get_alternative_by_name(alternative)
        gl_alternative_other = self._gl_switcher_other.get_alternative_by_name(alternative)

        # See if they are null
        # Abort if gl_alternative is null
        if gl_alternative and gl_alternative_other:
            success = (self._gl_switcher.set_alternative(gl_alternative) and
                       self._gl_switcher_other.set_alternative(gl_alternative_other))
            return success
        else:
            sys.stderr.write("Error: no alternative can be found for %s\n" % alternative)

        return False

    def print_current_alternative(self):
        alternatives = self._get_current_alternative()

        try:
            alternative = str(alternatives[0])
            alternative_other = str(alternatives[1])
        except IndexError:
            # No alternative exists
            return False

        if (alternative == alternative_other and
            'nvidia' in alternative):
            if 'prime' in alternative:
                sys.stdout.write('intel\n')
            else:
                sys.stdout.write('nvidia\n')
        else:
            sys.stdout.write('unknown\n')

        return True

    def _write_profile(self, profile):
        if profile == 'intel':
            nvidia_power = 'off'
        elif profile == 'nvidia':
            nvidia_power = 'on'
        else:
            return False

        # Write the settings to the file
        settings = open(self._power_profile_path, 'w')
        settings.write('%s\n' % nvidia_power)
        settings.close()


    def enable_profile(self, profile):
        current_driver = ''
        current_profile = ''
        alternatives = self._get_current_alternative()

        try:
            alternative = str(alternatives[0])
            alternative_other = str(alternatives[1])
        except IndexError:
            # No alternative exists
            sys.stderr.write('Error: the required alternatives do not exist\n')
            return False

        # The alternatives we are looking for are both provided by
        # an nvidia package
        if (alternative == alternative_other and
            'nvidia' in alternative):
            # The driver in use
            current_driver = alternative
            # The profile in use
            if 'prime' in alternative:
                current_profile = 'intel'
            else:
                current_profile = 'nvidia'
                # Check that the alternative for the intel profile
                # exists. If it doesn't, then the driver does not
                # support power management with PRIME
                alt_driver = '%s-prime' % current_driver
                if (not self._gl_switcher.get_alternative_by_name(alt_driver) or
                    not self._gl_switcher_other.get_alternative_by_name(alt_driver)):
                    sys.stderr.write('Error: %s does not support PRIME power management\n')
                    return False

            # Are alternatives in "auto" or in "manual" mode?
            status = self._get_current_alternative_status()

            # No need to do anything if we're already using the desired
            # profile, provided that alternatives are not in "auto" mode
            if profile == current_profile and 'auto' not in status:
                sys.stdout.write('Info: the %s profile is already in use\n' % (profile))
                return True
        elif (alternative != alternative_other and
            'nvidia' in alternative):
            # Something's wrong here
            # e.g. nvidia-319 on x86_64 and mesa on i386
            sys.stdout.write('Info: the alternatives do not match\n')
            # The first alternative shows the driver in use
            current_driver = alternative

            if 'prime' not in alternative:
                # Check that the alternative for the intel profile
                # exists. If it doesn't, then the driver does not
                # support power management with PRIME
                alt_driver = '%s-prime' % current_driver
                if (not self._gl_switcher.get_alternative_by_name(alt_driver) or
                    not self._gl_switcher_other.get_alternative_by_name(alt_driver)):
                    sys.stderr.write('Error: %s does not support PRIME power management\n')
                    return False
        else:
            sys.stderr.write('Error: alternatives are not set up properly\n')
            return False

        # Enable the desired alternative
        if profile == 'intel':
            target_driver = '%s-prime' % current_driver
        else:
            target_driver = current_driver.replace('-prime', '')

        if self.enable_alternative(target_driver):
            # Write the settings to the config file
            self._write_profile(profile)
            return True
        else:
            return False


def check_root():
    if not os.geteuid() == 0:
        sys.stderr.write("This operation requires root privileges\n")
        exit(1)

def handle_alternatives_error(mode):
    sys.stderr.write('Error: %s mode can\'t be enabled\n' % mode)
    exit(1)

def handle_query_error():
    sys.stderr.write("Error: no alternative can be found\n")
    exit(1)

def usage():
    sys.stderr.write("Usage: %s nvidia|intel|query\n" % (sys.argv[0]))

if __name__ == '__main__':
    try:
        arg = sys.argv[1]
    except IndexError:
        arg = None

    if len(sys.argv[1:]) != 1:
        usage()
        exit(1)

    switcher = Switcher()

    if arg == 'intel' or arg == 'nvidia':
        check_root()
        if not switcher.enable_profile(arg):
            handle_alternatives_error(arg)
    elif arg == 'query':
        if not switcher.print_current_alternative():
            handle_query_error()
    else:
        usage()
        sys.exit(1)

    exit(0)