Add country and region filtering in the xml parsing.

Move protocol filtering to the new xml filtering system.

3 files changed, 30 insertions, 9 deletions

diff --git a/mirrorselect/main.py b/mirrorselect/main.py
index 3f90694..5ba9116 100755
--- a/mirrorselect/main.py
+++ b/mirrorselect/main.py
@@ -271,6 +271,16 @@ class MirrorSelect(object):
 		group.add_option(
 			"-6", "--ipv6", action="store_true", default=False,
 			help="only use IPv6")
+		group.add_option(
+			"-c", "--country", action="store", default=None,
+			help="only use mirrors from the specified country "
+			"NOTE: Names with a space must be quoted "
+			"eg.:  -c 'South Korea'")
+		group.add_option(
+			"-R", "--region", action="store", default=None,
+			help="only use mirrors from the specified region"
+			"NOTE: Names with a space must be quoted"
+			"eg.:  -r 'North America'")
 
 		group = parser.add_option_group("Other options")
 		group.add_option(
diff --git a/mirrorselect/mirrorparser3.py b/mirrorselect/mirrorparser3.py
index e34f9c2..976ff17 100644
--- a/mirrorselect/mirrorparser3.py
+++ b/mirrorselect/mirrorparser3.py
@@ -35,7 +35,15 @@ MIRRORS_3_XML = 'http://www.gentoo.org/main/en/mirrors3.xml'
 MIRRORS_RSYNC_DATA = 'http://www.gentoo.org/main/en/mirrors-rsync-data.xml'
 
 class MirrorParser3:
-	def __init__(self):
+	def __init__(self, options=None):
+		self.filters = {}
+		for opt in ["country", "region"]:
+			value = getattr(options, opt)
+			if value is not None:
+				self.filters[opt] = value
+		for opt in ["ftp", "http"]:
+			if getattr(options, opt):
+				self.filters["proto"] = opt
 		self._reset()
 
 	def _reset(self):
@@ -51,7 +59,7 @@ class MirrorParser3:
 						name = e.text
 					if e.tag == 'uri':
 						uri = e.text
-						self._dict[uri] = {
+						data = {
 							"name": name,
 							"country": mirrorgroup.get("countryname"),
 							"region": mirrorgroup.get("region"),
@@ -59,6 +67,15 @@ class MirrorParser3:
 							"ipv6": e.get("ipv6"),
 							"proto": e.get("protocol"),
 							}
+						if len(self.filters):
+							good = True
+							for f in self.filters:
+								if data[f] != self.filters[f]:
+									good = False
+							if good:
+								self._dict[uri] = data
+						else:
+							self._dict[uri] = data
 
 	def tuples(self):
 		return [(url, args) for url, args in list(self._dict.items())]
diff --git a/mirrorselect/selectors.py b/mirrorselect/selectors.py
index 38ca07a..a603a47 100644
--- a/mirrorselect/selectors.py
+++ b/mirrorselect/selectors.py
@@ -58,19 +58,13 @@ class Extractor(object):
 
 	def __init__(self, list_url, options, output):
 		self.output = output
-		parser = MirrorParser3()
+		parser = MirrorParser3(options)
 		self.hosts = []
 
 		hosts = self.getlist(parser, list_url)
 		self.output.write('Extractor(): fetched mirrors.xml,'
 				' %s hosts before filtering\n' % len(hosts), 2)
 
-		if not options.rsync:
-			if options.ftp:
-				hosts = self.restrict_protocall('ftp', hosts)
-			if options.http:
-				hosts = self.restrict_protocall('http', hosts)
-
 		self.hosts = hosts
 
 	def restrict_protocall(self, prot, hosts):