reflector-2020.12.20.tar.xz

1 files changed, 83 insertions, 10 deletions

diff --git a/Reflector.py b/Reflector.py
index 7b6b7ab..26e5b99 100644
--- a/Reflector.py
+++ b/Reflector.py
@@ -167,6 +167,44 @@ def count_countries(mirrors):
     return countries
 
 
+def country_sort_key(priorities):
+    '''
+    Return a sort key function based on a list of country priorities.
+
+    Args:
+        priorities:
+            The list of countries in the order of priority. Any countries not in
+            the list will be sorted alphabetically after the countries in the
+            list. The countries may be specified by name or country code.
+
+    Returns:
+        A key function to pass to sort().
+    '''
+    priorities = [country.upper() for country in priorities]
+    try:
+        default_priority = priorities.index('*')
+    except ValueError:
+        default_priority = len(priorities)
+
+    def key_func(mirror):
+        country = mirror['country'].upper()
+        code = mirror['country_code'].upper()
+
+        try:
+            return (priorities.index(country), country)
+        except ValueError:
+            pass
+
+        try:
+            return (priorities.index(code), country)
+        except ValueError:
+            pass
+
+        return (default_priority, country)
+
+    return key_func
+
+
 # ------------------------ download timeout handling ------------------------- #
 
 class DownloadTimeout(Exception):
@@ -225,9 +263,31 @@ class DownloadTimer():
 
 # --------------------------------- Sorting ---------------------------------- #
 
-def sort(mirrors, by=None, **kwargs):  # pylint: disable=invalid-name
+def sort(mirrors, by=None, key=None, **kwargs):  # pylint: disable=invalid-name
     '''
     Sort mirrors by different criteria.
+
+    Args:
+        mirrors:
+            The iterable of mirrors to sort. This will be converted to a list.
+
+        by:
+            A mirrorstatus field by which to sort the mirrors, or one of the
+            following:
+
+            * age - Sort the mirrors by their last synchronization.
+            * rate - Sort the mirrors by download rate.
+
+        key:
+            A custom sorting function that accepts mirrors and returns a sort
+            key. If given, it will override the "by" parameter.
+
+        **kwargs:
+            Keyword arguments that are passed through to rate() when "by" is
+            "rate".
+
+    Returns:
+        The sorted mirrors as a list.
     '''
     # Ensure that "mirrors" is a list that can be sorted.
     if not isinstance(mirrors, list):
@@ -241,8 +301,12 @@ def sort(mirrors, by=None, **kwargs):  # pylint: disable=invalid-name
         mirrors = sorted(mirrors, key=lambda m: rates[m['url']], reverse=True)
 
     else:
+        if key is None:
+            def key(mir):
+                return mir[by]
         try:
-            mirrors.sort(key=lambda m: m[by])
+            print('sorting by', by, key)
+            mirrors.sort(key=key)
         except KeyError as err:
             raise MirrorStatusError('attempted to sort mirrors by unrecognized criterion: "{}"'.format(by)) from err
 
@@ -417,10 +481,11 @@ class MirrorStatusFilter():  # pylint: disable=too-many-instance-attributes
         mirrors = (m for m in mirrors if m['completion_pct'] >= self.min_completion_pct)
 
         # Filter by countries.
-        if self.countries:
+        countries = self.countries
+        if countries and '*' not in countries:
             mirrors = (
                 m for m in mirrors
-                if m['country'].upper() in self.countries or m['country_code'].upper() in self.countries
+                if m['country'].upper() in countries or m['country_code'].upper() in countries
             )
 
         # Filter by protocols.
@@ -648,9 +713,14 @@ class ListCountries(argparse.Action):
     def __call__(self, parser, namespace, values, option_string=None):
         ms = MirrorStatus(url=namespace.url)  # pylint: disable=invalid-name
         countries = ms.list_countries()
-        width = max(len(c) for c, cc in countries)
-        number = len(str(max(countries.values())))
-        fmt = '{{:{:d}s}} {{}} {{:{:d}d}}'.format(width, number)
+        headers = ('Country', 'Code', 'Count')
+        widths = [len(h) for h in headers]
+        widths[0] = max(widths[0], max(len(c) for c, cc in countries))
+        widths[2] = max(widths[2], len(str(max(countries.values()))))
+        fmt = '{{:{:d}s}} {{:>{:d}s}} {{:{:d}d}}'.format(*widths)
+        hdr_fmt = fmt.replace('d', 's')
+        print(hdr_fmt.format(*headers))
+        print(' '.join('-' * w for w in widths))
         for (ctry, count), nmbr in sorted(countries.items(), key=lambda x: x[0][0]):
             print(fmt.format(ctry, count, nmbr))
         sys.exit(0)
@@ -738,8 +808,8 @@ def add_arguments(parser):
     )
 
     filters.add_argument(
-        '-c', '--country', dest='countries', action='append', metavar='<country>',
-        help='Match one of the given countries (case-sensitive). Multiple countries may be selected using commas (e.g. "France,Germany") or by passing this option multiple times. Use "--list-countries" to see which are available.'
+        '-c', '--country', dest='countries', action='append', metavar='<country name or code>',
+        help='Restrict mirrors to selected countries. Countries may be given by name or country code, or a mix of both. The case is ignored. Multiple countries may be selected using commas (e.g. --country France,Germany) or by passing this option multiple times (e.g. -c fr -c de). Use "--list-countries" to display a table of available countries along with their country codes. When sorting by country, this option may also be used to sort by a preferred order instead of alphabetically. For example, to select mirrors from Sweden, Norway, Denmark and Finland, in that order, use the options "--country se,no,dk,fi --sort country". To set a preferred country sort order without filtering any countries. this option also recognizes the glob pattern "*", which will match any country. For example, to ensure that any mirrors from Sweden are at the top of the list and any mirrors from Denmark are at the bottom, with any other countries in between, use "--country \'se,*,dk\' --sort country". It is however important to note that when "*" is given along with other filter criteria, there is no guarantee that certain countries will be included in the results. For example, with the options "--country \'se,*,dk\' --sort country --latest 10", the latest 10 mirrors may all be from the United States. When the glob pattern is present, it only ensures that if certain countries are included in the results, they will be sorted in the requested order.'
     )
 
     filters.add_argument(
@@ -883,7 +953,10 @@ def process_options(options, mirrorstatus=None, mirrors=None):
         mirrors = itertools.islice(mirrors, options.fastest)
 
     if options.sort and not (options.sort == 'rate' and options.fastest):
-        mirrors = mirrorstatus.sort(mirrors, by=options.sort)
+        if options.sort == 'country' and options.countries:
+            mirrors = mirrorstatus.sort(mirrors, key=country_sort_key(options.countries))
+        else:
+            mirrors = mirrorstatus.sort(mirrors, by=options.sort)
 
     if options.number:
         mirrors = list(mirrors)[:options.number]