Windows fix + Sorting based on list

This fix introduces changes so that development can be done (and tested) on other platforms than Linux. This is a convenience fix and shouldn't break anything (simply a few Linux-specific imports that have moved into the functions where they are used).

This commit also introduces sorting based on a list of priorities (where the default will be last if not matched).

1 files changed, 62 insertions, 6 deletions

diff --git a/archinstall/lib/mirrors.py b/archinstall/lib/mirrors.py
index fad3b129..44c73132 100644
--- a/archinstall/lib/mirrors.py
+++ b/archinstall/lib/mirrors.py
@@ -4,8 +4,53 @@ import urllib.request
 from .general import *
 from .output import log
 
+def sort_mirrorlist(raw_data :bytes, sort_order=["https", "http"]) -> bytes:
+	"""
+	This function can sort /etc/pacman.d/mirrorlist according to the
+	mirror's URL prefix. By default places HTTPS before HTTP but it also
+	preserves the country/rank-order.
+
+	This assumes /etc/pacman.d/mirrorlist looks like the following:
+
+	## Comment
+	Server = url
+
+	or
+
+	## Comment
+	#Server = url
+
+	But the Comments need to start with double-hashmarks to be distringuished
+	from server url definitions (commented or uncommented).
+	"""
+	comments_and_whitespaces = b""
+
+	categories = {key: [] for key in sort_order+["Unknown"]}
+	for line in raw_data.split(b"\n"):
+		if line[0:2] in (b'##', b''):
+			comments_and_whitespaces += line + b'\n'
+		elif line[:6].lower() == b'server' or line[:7].lower() == b'#server':
+			opening, url = line.split(b'=', 1)
+			opening, url = opening.strip(), url.strip()
+			if (category := url.split(b'://',1)[0].decode('UTF-8')) in categories:
+				categories[category].append(comments_and_whitespaces)
+				categories[category].append(opening+b' = '+url+b'\n')
+			else:
+				categories["Unknown"].append(comments_and_whitespaces)
+				categories["Unknown"].append(opening+b' = '+url+b'\n')
+
+			comments_and_whitespaces = b""
 
-def filter_mirrors_by_region(regions, destination='/etc/pacman.d/mirrorlist', *args, **kwargs):
+
+	new_raw_data = b''
+	for category in sort_order+["Unknown"]:
+		for line in categories[category]:
+			new_raw_data += line
+
+	return new_raw_data
+
+
+def filter_mirrors_by_region(regions, destination='/etc/pacman.d/mirrorlist', sort_order=["https", "http"], *args, **kwargs):
 	"""
 	This function will change the active mirrors on the live medium by
 	filtering which regions are active based on `regions`.
@@ -18,10 +63,17 @@ def filter_mirrors_by_region(regions, destination='/etc/pacman.d/mirrorlist', *a
 		region_list.append(f'country={region}')
 	response = urllib.request.urlopen(urllib.request.Request(f"https://archlinux.org/mirrorlist/?{'&'.join(region_list)}&protocol=https&protocol=http&ip_version=4&ip_version=6&use_mirror_status=on'", headers={'User-Agent': 'ArchInstall'}))
 	new_list = response.read().replace(b"#Server", b"Server")
-	with open(destination, "wb") as mirrorlist:
-		mirrorlist.write(new_list)
 
-	return True
+	if sort_order:
+		new_list = sort_mirrorlist(new_list, sort_order=sort_order)
+
+	if destination:
+		with open(destination, "wb") as mirrorlist:
+			mirrorlist.write(new_list)
+
+		return True
+	else:
+		return new_list.decode('UTF-8')
 
 
 def add_custom_mirrors(mirrors: list, *args, **kwargs):
@@ -78,7 +130,7 @@ def re_rank_mirrors(top=10, *positionals, **kwargs):
 	return False
 
 
-def list_mirrors():
+def list_mirrors(sort_order=["https", "http"]):
 	url = "https://archlinux.org/mirrorlist/?protocol=https&protocol=http&ip_version=4&ip_version=6&use_mirror_status=on"
 	regions = {}
 
@@ -88,8 +140,12 @@ def list_mirrors():
 		log(f'Could not fetch an active mirror-list: {err}', level=logging.WARNING, fg="yellow")
 		return regions
 
+	mirrorlist = response.read()
+	if sort_order:
+		mirrorlist = sort_mirrorlist(mirrorlist, sort_order=sort_order)
+
 	region = 'Unknown region'
-	for line in response.readlines():
+	for line in mirrorlist.split(b'\n'):
 		if len(line.strip()) == 0:
 			continue