Move the filtering code back into Extractor class

2 files changed, 35 insertions, 40 deletions

diff --git a/mirrorselect/mirrorparser3.py b/mirrorselect/mirrorparser3.py
index 6a11855..0db0edb 100644
--- a/mirrorselect/mirrorparser3.py
+++ b/mirrorselect/mirrorparser3.py
@@ -36,14 +36,6 @@ MIRRORS_RSYNC_DATA = 'http://www.gentoo.org/main/en/mirrors-rsync-data.xml'
 
 class MirrorParser3:
 	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):
@@ -59,7 +51,7 @@ class MirrorParser3:
 						name = e.text
 					if e.tag == 'uri':
 						uri = e.text
-						data = {
+						self._dict[uri] = {
 							"name": name,
 							"country": mirrorgroup.get("countryname"),
 							"region": mirrorgroup.get("region"),
@@ -67,15 +59,6 @@ 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 a603a47..df9dcd7 100644
--- a/mirrorselect/selectors.py
+++ b/mirrorselect/selectors.py
@@ -58,32 +58,44 @@ class Extractor(object):
 
 	def __init__(self, list_url, options, output):
 		self.output = output
-		parser = MirrorParser3(options)
+		filters = {}
+		for opt in ["country", "region"]:
+			value = getattr(options, opt)
+			if value is not None:
+				filters[opt] = value
+				self.output.print_info('Limiting test to %s = %s hosts. \n'
+					%(opt, value))
+		for opt in ["ftp", "http"]:
+			if getattr(options, opt):
+				filters["proto"] = opt
+				self.output.print_info('Limiting test to %s hosts. \n' % opt )
+		parser = MirrorParser3()
 		self.hosts = []
 
-		hosts = self.getlist(parser, list_url)
-		self.output.write('Extractor(): fetched mirrors.xml,'
-				' %s hosts before filtering\n' % len(hosts), 2)
-
-		self.hosts = hosts
-
-	def restrict_protocall(self, prot, hosts):
-		"""
-		Removes hosts that are not of the specified type.
-		"prot" must always be exactly 'http' or 'ftp'.
-		"""
-		myhosts = []
+		self.unfiltered_hosts = self.getlist(parser, list_url)
 
-		self.output.print_info('Limiting test to %s hosts. ' % prot )
+		self.hosts = self.filter_hosts(filters, self.unfiltered_hosts)
 
-		for host in hosts:
-			if host[0].startswith(prot):
-				myhosts.append(host)
+		self.output.write('Extractor(): fetched mirrors.xml,'
+				' %s hosts after filtering\n' % len(self.hosts), 2)
 
-		self.output.write('%s of %s removed.\n' % (len(hosts) - len(myhosts),
-			len(hosts)) )
 
-		return myhosts
+	@staticmethod
+	def filter_hosts(filters, hosts):
+		"""Filter the hosts to the criteria passed in
+		Return the filtered list
+		"""
+		if not len(filters):
+			return hosts
+		filtered = []
+		for uri, data in hosts:
+			good = True
+			for f in filters:
+				if data[f] != filters[f]:
+					good = False
+			if good:
+				filtered.append((uri, data))
+		return filtered
 
 
 	def getlist(self, parser, url):
@@ -102,8 +114,8 @@ class Extractor(object):
 			pass
 
 		if len(parser.tuples()) == 0:
-			self.output.print_err('Could not get mirror list. Check your internet'
-					' connection.')
+			self.output.print_err('Could not get mirror list. '
+				'Check your internet connection.')
 
 		self.output.write(' Got %d mirrors.\n' % len(parser.tuples()))