Reload selected repositories

reload_repository.py

#!/usr/bin/env python
# Copyright (C) 2012 Glenn Washburn

# 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.

# TODO:
# * Write a more informative progress class that says which repo is being
#   updated.
# * Add a gtk gui for easy selection/deselection

# libapt-pkg documentation at:
# file:///usr/share/doc/libapt-pkg-doc/html/index.html

# A simpler form of what we do here:
# apt-get update -o Dir::Etc::sourcelist="sources.list.d/other.list"
#                -o Dir::Etc::sourceparts="-" -o APT::Get::List-Cleanup="0"

import sys
import os
import getopt

import apt
import apt_pkg


class PrintableSourceList(apt_pkg.SourceList):
    def __str__(self):
        return '\n'.join(self.all_urls())
    
    def all_urls(self):
        for item in self.list:
            for index in item.index_files:
                yield index.describe


class CustomAcquireProgress:
    pass

def acquire_progress_factory(ui='text'):
    mod = getattr(__import__("apt.progress", fromlist=[ui]), ui)
    if ui == 'gtk2':
        base = mod.GAcquireProgress
    else:
        base = mod.AcquireProgress
    
    mixin = CustomAcquireProgress
    
    class AcquireProgress(mixin, base):
        pass
    
    return AcquireProgress

def get_repositories(repos):
    # Create a temporary sources.list file to load the SourceList object from.
    # We can currently only create a SourceList objects by reading the
    # sources.list or its partial directory defined in the config.  So we
    # set the setting to a tmp file in which we stuff the desired repo entries.
    path = '/tmp/reloadsources.list.%d'%os.getpid()
    fsrclist = open(path, 'w')
    
    # May search the system sources.list.d directory later. Get it now.
    srcparts = apt_pkg.config.get("Dir::Etc::sourceparts")
    if srcparts[0] != "/":
        srcparts = os.path.join("/etc/apt", srcparts)
    
    for repo in repos:
        # If repo is a file path, then assume is a sources.list fragment
        if os.path.isfile(repo):
            print >> fsrclist, open(repo).read()
        # Might also be just the base name of a sources.list.d fragment
        elif os.path.isfile(os.path.join(srcparts, "%s.list"%repo)):
            print >> fsrclist, open(os.path.join(srcparts, "%s.list"%repo)).read()
        else:
            # Assume is a line in sources.list
            print >> fsrclist, repo
    fsrclist.close()
    
    # Tell apt that it should use the new sources.list.
    apt_pkg.config.set("Dir::Etc::sourcelist", path)
    apt_pkg.config.set("Dir::Etc::sourceparts", "-")
    
    #~ print apt_pkg.config.get("Dir::Etc::sourcelist")
    #~ print open(apt_pkg.config.get("Dir::Etc::sourcelist")).read()
    
    srclist = PrintableSourceList()
    try:
        srclist.read_main_list()
    except SystemError as err:
        print >> sys.stderr, err
    finally:
        os.unlink(path)
    
    return srclist

def main(argv):
    
    verbose = False
    ui = 'text'
    
    opts, args = getopt.gnu_getopt(argv, "vu:")
    for opt in opts:
        if opt[0] == '-v':
            verbose = True
        if opt[0] == '-u':
            ui = opt[1]
    
    # ensure that other lists are not cleaned up (removed) because we only
    # want to update the ones we specify.
    apt_pkg.config.set('APT::Get::List-Cleanup', '0')
    
    # Get desired repos to update as a sourcelist
    srclist = get_repositories(args)
    
    if verbose:
        print srclist
    
    try:
        # Lock pkg dir so no one else is updating or trying to read
        with apt_pkg.SystemLock():
            cc = apt_pkg.Cache(None)
            
            AcquireProgress = acquire_progress_factory(ui)
            cc.update(AcquireProgress(), srclist)
    except SystemError as err:
        print >> sys.stderr, err


if __name__ == "__main__":
    main(sys.argv[1:])