Merge pull request #70 from jblakeman/release-2.1.0

Release 2.1.0

apt_select/__init__.py

@@ -1 +1 @@
-__version__ = '2.0.0'
+__version__ = '2.1.0'

apt_select/__main__.py

@@ -8,7 +8,7 @@
 from os import getcwd
 from apt_select.arguments import get_args, DEFAULT_COUNTRY
 from apt_select.mirrors import Mirrors
-from apt_select.apt_system import AptSources, SourcesFileError
+from apt_select.apt import System, Sources, SourcesFileError
 from apt_select.utils import DEFAULT_REQUEST_HEADERS
 
 # Support input for Python 2 and 3
@@ -130,14 +130,18 @@ def yes_or_no(query):
 
 def apt_select():
     """Run apt-select: Ubuntu archive mirror reporting tool"""
-    args = set_args()
+
+    try:
+        system = System()
+    except OSError as err:
+        exit("Error setting system information:\n\t%s" % err)
+
     try:
-        sources = AptSources()
-    except ValueError as err:
-        exit("Error setting system information: %s" % err)
+        sources = Sources(system.codename)
     except SourcesFileError as err:
-        exit("Error with current apt sources: %s" % err)
+        exit("Error with current apt sources:\n\t%s" % err)
 
+    args = set_args()
     mirrors_loc = "mirrors.ubuntu.com"
     mirrors_url = "http://%s/%s.txt" % (mirrors_loc, args.country.upper())
     mirrors_list = get_mirrors(mirrors_url, args.country)
@@ -156,7 +160,11 @@ def apt_select():
             # Mirrors needs a limit to stop launching threads
             archives.status_num = args.top_number
             stderr.write("Looking up %d status(es)\n" % args.top_number)
-            archives.lookup_statuses(args.min_status)
+            archives.lookup_statuses(
+                system.codename.capitalize(),
+                system.arch,
+                args.min_status
+            )
 
         if args.top_number > 1:
             stderr.write('\n')
@@ -165,7 +173,7 @@ def apt_select():
         archives.top_list = archives.ranked[:args.top_number]
 
     sources.set_current_archives()
-    current_url = sources.urls[0]
+    current_url = sources.urls['current']
     if archives.urls.get(current_url):
         archives.urls[current_url]['Host'] += " (current)"
 
@@ -207,15 +215,15 @@ def set_hostname_len(url, i):
             }))
 
     work_dir = getcwd()
-    if work_dir == sources.directory[0:-1]:
+    if work_dir == sources.DIRECTORY[0:-1]:
         query = (
             "'%(dir)s' is the current directory.\n"
             "Generating a new '%(apt)s' file will "
             "overwrite the current file.\n"
             "You should copy or backup '%(apt)s' before replacing it.\n"
             "Continue?\n[yes|no] " % {
-                'dir': sources.directory,
-                'apt': sources.apt_file
+                'dir': sources.DIRECTORY,
+                'apt': sources.APT_FILE
             }
         )
         yes_or_no(query)

apt_select/apt_system.py → apt_select/apt.py

@@ -4,42 +4,80 @@
 from os import path
 from apt_select.utils import utf8_decode
 
+SUPPORTED_KERNEL = 'Linux'
+SUPPORTED_DISTRIBUTION_TYPE = 'Ubuntu'
+
+UNAME = 'uname'
+KERNEL_COMMAND = (UNAME, '-s')
+MACHINE_COMMAND = (UNAME, '-m')
+RELEASE_COMMAND = ('lsb_release', '-ics')
+RELEASE_FILE = '/etc/lsb-release'
+
 LAUNCHPAD_ARCH_32 = 'i386'
 LAUNCHPAD_ARCH_64 = 'amd64'
 LAUNCHPAD_ARCHES = frozenset([
     LAUNCHPAD_ARCH_32,
     LAUNCHPAD_ARCH_64
 ])
 
-class AptSystem(object):
-    """System information for use in apt related operations"""
-
-    @staticmethod
-    def get_release():
-        """Call system for Ubuntu release information"""
-        return [utf8_decode(s).strip()
-                for s in check_output(["lsb_release", "-ics"]).split()]
-
-    @staticmethod
-    def get_arch():
-        """Return architecture information in Launchpad format"""
-        if utf8_decode(check_output(["uname", "-m"]).strip()) == 'x86_64':
-            return LAUNCHPAD_ARCH_64
-
-        return LAUNCHPAD_ARCH_32
 
-    _not_ubuntu = "Must be an Ubuntu OS"
-    try:
-        dist, codename = get_release.__func__()
-    except OSError:
-        raise ValueError("%s\n%s" % _not_ubuntu)
-    else:
-        codename = codename.capitalize()
+class System(object):
+    """System information for use in apt related operations"""
 
-    if dist != 'Ubuntu':
-        raise ValueError(_not_ubuntu)
+    def __init__(self):
+        _kernel = utf8_decode(check_output(KERNEL_COMMAND)).strip()
+        if _kernel != SUPPORTED_KERNEL:
+            raise OSError(
+                "Invalid kernel found: %s. Expected %s." % (
+                    _kernel,
+                    SUPPORTED_KERNEL,
+                )
+            )
 
-    arch = get_arch.__func__()
+        try:
+            self.dist, self.codename = tuple(
+                utf8_decode(s).strip()
+                for s in check_output(RELEASE_COMMAND).split()
+            )
+        except OSError:
+            # Fall back to using lsb-release info file if lsb_release command
+            # is not available. e.g. Ubuntu minimal (core, docker image).
+            try:
+                with open(RELEASE_FILE, 'rU') as release_file:
+                    try:
+                        lsb_info = dict(
+                            line.strip().split('=')
+                            for line in release_file.readlines()
+                        )
+                    except ValueError:
+                        raise OSError(
+                            "Unexpected release file format found in %s." % RELEASE_FILE
+                        )
+
+                    try:
+                        self.dist = lsb_info['DISTRIB_ID']
+                        self.codename = lsb_info['DISTRIB_CODENAME']
+                    except KeyError:
+                        raise OSError(
+                            "Expected distribution keys missing from %s." % RELEASE_FILE
+                        )
+
+            except (IOError, OSError):
+                raise OSError((
+                    "Unable to determine system distribution. "
+                    "%s is required." % SUPPORTED_DISTRIBUTION_TYPE
+                ))
+
+        if self.dist != SUPPORTED_DISTRIBUTION_TYPE:
+            raise OSError(
+                "%s distributions are not supported. %s is required." % (
+                    self.dist, SUPPORTED_DISTRIBUTION_TYPE
+                )
+            )
+
+        self.arch = LAUNCHPAD_ARCH_32
+        if utf8_decode(check_output(MACHINE_COMMAND).strip()) == 'x86_64':
+            self.arch = LAUNCHPAD_ARCH_64
 
 
 class SourcesFileError(Exception):
@@ -51,19 +89,21 @@ class SourcesFileError(Exception):
     pass
 
 
-class AptSources(AptSystem):
+class Sources(object):
     """Class for apt configuration files"""
 
     DEB_SCHEMES = frozenset(['deb', 'deb-src'])
     PROTOCOLS = frozenset(['http', 'ftp', 'https'])
 
-    def __init__(self):
-        self.directory = '/etc/apt/'
-        self.apt_file = 'sources.list'
-        self._config_path = self.directory + self.apt_file
-        if not path.isfile(self._config_path):
+    DIRECTORY = '/etc/apt/'
+    LIST_FILE = 'sources.list'
+    _CONFIG_PATH = DIRECTORY + LIST_FILE
+
+    def __init__(self, codename):
+        self._codename = codename.lower()
+        if not path.isfile(self._CONFIG_PATH):
             raise SourcesFileError((
-                "%s must exist as file" % self._config_path
+                "%s must exist as file" % self._CONFIG_PATH
             ))
 
         self._required_component = "main"
@@ -76,7 +116,7 @@ def __set_sources_lines(self):
         """Read system config file and store the lines in memory for parsing
            and generation of new config file"""
         try:
-            with open(self._config_path, 'r') as f:
+            with open(self._CONFIG_PATH, 'r') as f:
                 self._lines = f.readlines()
         except IOError as err:
             raise SourcesFileError((
@@ -94,22 +134,16 @@ def __confirm_apt_source_uri(self, uri):
     def __get_current_archives(self):
         """Parse through all lines of the system apt file to find current
            mirror urls"""
-        urls = []
-        cname = self.codename.lower()
+        urls = {}
         for line in self._lines:
             fields = line.split()
             if self.__confirm_apt_source_uri(fields):
                 if (not urls and
-                        (cname in fields[2]) and
+                        (self._codename in fields[2]) and
                         (fields[3] == self._required_component)):
-                    urls.append(fields[1])
-                    continue
-                elif (urls and
-                        (fields[2] == '%s-security' % cname) and
-                        # Mirror urls should be unique as they'll be
-                        # used in a global search and replace
-                        (urls[0] != fields[1])):
-                    urls.append(fields[1])
+                    urls['current'] = fields[1]
+                elif urls and (fields[2] == '%s-security' % self._codename):
+                    urls['security'] = fields[1]
                     break
 
         return urls
@@ -126,7 +160,7 @@ def set_current_archives(self):
         if not urls:
             raise SourcesFileError((
                 "Error finding current %s URI in %s\n%s\n" %
-                (self._required_component, self._config_path,
+                (self._required_component, self._CONFIG_PATH,
                  self.skip_gen_msg)
             ))
 
@@ -135,13 +169,13 @@ def set_current_archives(self):
     def __set_config_lines(self, new_mirror):
         """Replace all instances of the current urls with the new mirror"""
         self._lines = ''.join(self._lines)
-        for url in self.urls:
+        for url in self.urls.values():
             self._lines = self._lines.replace(url, new_mirror)
 
     def generate_new_config(self, work_dir, new_mirror):
         """Write new configuration file to current working directory"""
         self.__set_config_lines(new_mirror)
-        self.new_file_path = work_dir.rstrip('/') + '/' + self.apt_file
+        self.new_file_path = work_dir.rstrip('/') + '/' + self.LIST_FILE
         try:
             with open(self.new_file_path, 'w') as f:
                 f.write(self._lines)

apt_select/mirrors.py

@@ -8,7 +8,6 @@
                     gethostbyname, error, timeout, gaierror)
 from time import time
 from apt_select.utils import progress_msg, get_text, URLGetTextError
-from apt_select.apt_system import AptSystem
 try:
     from urlparse import urlparse
 except ImportError:
@@ -159,7 +158,7 @@ def get_rtts(self):
             self.urls, key=lambda x: self.urls[x]["Latency"]
         )
 
-    def __queue_lookups(self, data_queue):
+    def __queue_lookups(self, codename, arch, data_queue):
         """Queue threads for data retrieval from launchpad.net
 
            Returns number of threads started to fulfill number of
@@ -172,7 +171,13 @@ def __queue_lookups(self, data_queue):
                 pass
             else:
                 thread = Thread(
-                    target=_LaunchData(url, launch_url, data_queue).get_info
+                    target=_LaunchData(
+                        url,
+                        launch_url,
+                        codename,
+                        arch,
+                        data_queue
+                    ).get_info
                 )
                 thread.daemon = True
                 thread.start()
@@ -187,11 +192,11 @@ def __queue_lookups(self, data_queue):
 
         return num_threads
 
-    def lookup_statuses(self, min_status):
+    def lookup_statuses(self, codename, arch, min_status):
         """Scrape statuses/info in from launchpad.net mirror pages"""
         while (self.got["data"] < self.status_num) and self.ranked:
             data_queue = Queue()
-            num_threads = self.__queue_lookups(data_queue)
+            num_threads = self.__queue_lookups(codename, arch, data_queue)
             if num_threads == 0:
                 break
             # Get output of all started thread methods from queue
@@ -220,9 +225,7 @@ def lookup_statuses(self, min_status):
                     break
 
             # Reorder by latency as queue returns vary building final list
-            self.top_list = sorted(
-                self.top_list, key=lambda x: self.urls[x]["Latency"]
-            )
+            self.top_list.sort(key=lambda x: self.urls[x]["Latency"])
 
             data_queue.join()
 
@@ -268,10 +271,12 @@ def min_rtt(self):
         self._trip_queue.put((self._url, min(rtts)))
 
 
-class _LaunchData(AptSystem):
-    def __init__(self, url, launch_url, data_queue):
+class _LaunchData(object):
+    def __init__(self, url, launch_url, codename, arch, data_queue):
         self._url = url
         self._launch_url = launch_url
+        self._codename = codename
+        self._arch = arch
         self._data_queue = data_queue
 
     def __parse_mirror_html(self, launch_html):
@@ -284,8 +289,8 @@ def __parse_mirror_html(self, launch_html):
                 # series name and machine architecture
                 for tr in line.find('tbody').find_all('tr'):
                     arches = [x.get_text() for x in tr.find_all('td')]
-                    if (self.codename in arches[0] and
-                            arches[1] == self.arch):
+                    if (self._codename in arches[0] and
+                            arches[1] == self._arch):
                         info.update({"Status": arches[2]})
             else:
                 # "Speed" lives in a dl, and we use the key -> value as such