Merge pull request #65 from jblakeman/devel

Devel to release

README.rst

@@ -6,7 +6,7 @@ Find a fast, up-to-date Ubuntu Archive Mirror.
 Features
 --------
 
-* Tests latency to mirrors in `mirrors.txt <http://mirrors.ubuntu.com/mirrors.txt>`_.
+* Tests latency to mirrors in a given country's mirror list at `mirrors.ubuntu.com <http://mirrors.ubuntu.com>`_.
     - 3 requests are sent to each mirror, minumum round trip time being used for rank.
 
 * Reports latency, status, and bandwidth capacity of the fastest mirrors in a ranked list.
@@ -40,17 +40,23 @@ Invocation
 ::
 
     $ apt-select --help
-    usage: apt-select [-h] [-t NUMBER] [-m STATUS | -p] [-c | -l]
+    usage: apt-select [-h] [-C [COUNTRY]] [-t [NUMBER]] [-m [STATUS] | -p]
+                      [-c | -l]
 
     Find the fastest Ubuntu apt mirrors.
     Generate new sources.list file.
 
     optional arguments:
       -h, --help            show this help message and exit
-      -t NUMBER, --top-number NUMBER
+      -C [COUNTRY], --country [COUNTRY]
+                            specify a country to test its list of mirrors
+                            used to match country list file names found at mirrors.ubuntu.com
+                            COUNTRY should follow ISO 3166-1 alpha-2 format
+                            default: US
+      -t [NUMBER], --top-number [NUMBER]
                             specify number of mirrors to return
                             default: 1
-      -m STATUS, --min-status STATUS
+      -m [STATUS], --min-status [STATUS]
                             return mirrors with minimum status
                             choices:
                                up-to-date
@@ -64,16 +70,20 @@ Invocation
       -c, --choose          choose mirror from a list
                             requires -t/--top-num NUMBER where NUMBER > 1
       -l, --list            print list of mirrors only, don't generate file
-                        cannot be used with -c/--choose
+                            cannot be used with -c/--choose
 
 Examples
 --------
 
+Get the top mirror from the United Kingdom to generate a new `sources.list`:::
+
+    apt-select --country GB
+
 Choose from the top 3 mirrors, including those last updated a week ago:::
 
     apt-select -c -t 3 -m one-week-behind
 
-Find the top 10 mirrors, output latency info only, and don't generate new config file:::
+Find the top 10 mirrors, output latency info only, and don't generate new `sources.list`:::
 
     apt-select -t 10 -p -l
 
@@ -85,11 +95,11 @@ After new sources.list is generated in current working directory, backup and rep
 Supported URI Types
 -------------------
 
-Currently, `http` and `ftp` are supported.
+Currently, `http`, `https` and `ftp` are supported.
 
 `/etc/apt/sources.list` should contain sources in the following format:::
 
-    [deb|deb-src] [http|ftp]://mirror.example.com/path [component1] [component2] [...]
+    [deb|deb-src] [http|https|ftp]://mirror.example.com/path [component1] [component2] [...]
 
 Dependencies
 ------------

apt_select/__main__.py

@@ -1,12 +1,15 @@
 #!/usr/bin/env python
 """Main apt-select script"""
 
+import requests
+import re
+
 from sys import exit, stderr, version_info
 from os import getcwd
-from apt_select.arguments import get_args
-from apt_select.utils import get_text, URLGetTextError
+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.utils import DEFAULT_REQUEST_HEADERS
 
 # Support input for Python 2 and 3
 get_input = input
@@ -31,19 +34,32 @@ def set_args():
             "where NUMBER is greater than 1."
         ))
 
+    if not args.country:
+        stderr.write('WARNING: no country code provided. defaulting to US.\n')
+        args.country = DEFAULT_COUNTRY
+    elif not re.match(r'^[a-zA-Z]{2}$', args.country):
+        exit((
+            "Invalid country. %s is not in ISO 3166-1 alpha-2 "
+            "format" % args.country
+        ))
+
     return args
 
 
-def get_mirrors(mirrors_url):
+def get_mirrors(mirrors_url, country):
     """Fetch list of Ubuntu mirrors"""
     stderr.write("Getting list of mirrors...")
-    try:
-        mirrors_list = get_text(mirrors_url)
-    except URLGetTextError as err:
-        exit("Error getting list from %s:\n\t%s" % (mirrors_url, err))
+    response = requests.get(mirrors_url, headers=DEFAULT_REQUEST_HEADERS)
+    if response.status_code == requests.codes.NOT_FOUND:
+        exit(
+            "The mirror list for country: %s was not found at %s" % (
+                country, mirrors_url
+            )
+        )
+
     stderr.write("done.\n")
 
-    return mirrors_list.splitlines()
+    return response.text.splitlines()
 
 
 def print_status(info, rank):
@@ -123,8 +139,8 @@ def apt_select():
         exit("Error with current apt sources: %s" % err)
 
     mirrors_loc = "mirrors.ubuntu.com"
-    mirrors_url = "http://%s/mirrors.txt" % mirrors_loc
-    mirrors_list = get_mirrors(mirrors_url)
+    mirrors_url = "http://%s/%s.txt" % (mirrors_loc, args.country.upper())
+    mirrors_list = get_mirrors(mirrors_url, args.country)
 
     archives = Mirrors(mirrors_list, args.ping_only, args.min_status)
     archives.get_rtts()

apt_select/apt_system.py

@@ -55,7 +55,7 @@ class AptSources(AptSystem):
     """Class for apt configuration files"""
 
     DEB_SCHEMES = frozenset(['deb', 'deb-src'])
-    PROTOCOLS = frozenset(['http', 'ftp'])
+    PROTOCOLS = frozenset(['http', 'ftp', 'https'])
 
     def __init__(self):
         self.directory = '/etc/apt/'

apt_select/arguments.py

@@ -3,6 +3,15 @@
 
 from argparse import ArgumentParser, RawTextHelpFormatter
 
+DEFAULT_COUNTRY = 'US'
+DEFAULT_NUMBER = 1
+STATUS_ARGS = (
+    "up-to-date",
+    "one-day-behind",
+    "two-days-behind",
+    "one-week-behind",
+    "unknown"
+)
 
 def get_args():
     """Get parsed command line arguments"""
@@ -13,6 +22,19 @@ def get_args():
         ),
         formatter_class=RawTextHelpFormatter
     )
+    parser.add_argument(
+        '-C',
+        '--country',
+        nargs='?',
+        type=str,
+        help=(
+            "specify a country to test its list of mirrors\n"
+            "used to match country list file names found at mirrors.ubuntu.com\n"
+            "COUNTRY should follow ISO 3166-1 alpha-2 format\n"
+            "default: %s" % DEFAULT_COUNTRY
+        ),
+        metavar='COUNTRY'
+    )
     parser.add_argument(
         '-t',
         '--top-number',
@@ -22,23 +44,16 @@ def get_args():
             "specify number of mirrors to return\n"
             "default: 1\n"
         ),
-        const=1,
-        default=1,
+        const=DEFAULT_NUMBER,
+        default=DEFAULT_NUMBER,
         metavar='NUMBER'
     )
-    status_args = (
-        "up-to-date",
-        "one-day-behind",
-        "two-days-behind",
-        "one-week-behind",
-        "unknown"
-    )
     test_group = parser.add_mutually_exclusive_group(required=False)
     test_group.add_argument(
         '-m',
         '--min-status',
         nargs='?',
-        choices=status_args,
+        choices=STATUS_ARGS,
         type=str,
         help=(
             "return mirrors with minimum status\n"
@@ -49,15 +64,15 @@ def get_args():
             "   %(week)s\n"
             "   %(unknown)s\n"
             "default: %(up)s\n" % {
-                'up': status_args[0],
-                'day': status_args[1],
-                'two_day': status_args[2],
-                'week': status_args[3],
-                'unknown': status_args[4]
+                'up': STATUS_ARGS[0],
+                'day': STATUS_ARGS[1],
+                'two_day': STATUS_ARGS[2],
+                'week': STATUS_ARGS[3],
+                'unknown': STATUS_ARGS[4]
             }
         ),
-        const=status_args[0],
-        default=status_args[0],
+        const=STATUS_ARGS[0],
+        default=STATUS_ARGS[0],
         metavar='STATUS'
     )
     test_group.add_argument(

apt_select/utils.py

@@ -2,8 +2,13 @@
 """Collection of module neutral utility functions"""
 
 from sys import stderr
+
 import requests
 
+DEFAULT_REQUEST_HEADERS = {
+    'User-Agent': 'apt-select'
+}
+
 
 def utf8_decode(encoded):
     return encoded.decode('utf-8')
@@ -17,7 +22,7 @@ class URLGetTextError(Exception):
 def get_text(url):
     """Return text from GET request response content"""
     try:
-        text = requests.get(url).text
+        text = requests.get(url, headers=DEFAULT_REQUEST_HEADERS).text
     except requests.HTTPError as err:
         raise URLGetTextError(err)