From 82d5ebb630ca5471920cf1deb9785bd45bf1e856 Mon Sep 17 00:00:00 2001
From: Felix Yan <felixonmars@archlinux.org>
Date: Mon, 21 Sep 2020 18:49:23 +0800
Subject: [PATCH] apt: Parse Packages and cache it to improve performance

It's a very noticable 20~40s to 1s improvement here.
---
 nvchecker_source/apt.py | 51 ++++++++++++++++++++++++++++++-----------
 1 file changed, 37 insertions(+), 14 deletions(-)

diff --git a/nvchecker_source/apt.py b/nvchecker_source/apt.py
index 4a6a489..e04ef59 100644
--- a/nvchecker_source/apt.py
+++ b/nvchecker_source/apt.py
@@ -4,6 +4,8 @@
 from __future__ import annotations
 
 import asyncio
+from io import StringIO
+from typing import Dict, Tuple
 
 from nvchecker.api import (
   session, GetVersionError,
@@ -33,6 +35,30 @@ async def get_url(url: str) -> str:
     None, _decompress_data,
     url, data)
 
+async def parse_packages(key: Tuple[AsyncCache, str]) -> Tuple[Dict[str, str], Dict[str, str]]:
+  cache, url = key
+  apt_packages = await cache.get(url, get_url) # type: ignore
+
+  pkg_map = {}
+  srcpkg_map = {}
+
+  pkg = None
+  srcpkg = None
+  for line in apt_packages.split('\n'):
+    if line.startswith("Package: "):
+      pkg = line[9:]
+    elif line.startswith("Source: "):
+      srcpkg = line[8:]
+    elif line.startswith("Version: "):
+      version = line[9:]
+      if pkg is not None:
+        pkg_map[pkg] = version
+      if srcpkg is not None:
+        srcpkg_map[srcpkg] = version
+      pkg = srcpkg = None
+
+  return pkg_map, srcpkg_map
+
 async def get_version(
   name: str, conf: Entry, *,
   cache: AsyncCache, keymanager: KeyManager,
@@ -60,19 +86,16 @@ async def get_version(
   else:
     raise GetVersionError('Packages file not found in APT repository')
 
-  apt_packages = await cache.get(
-    APT_PACKAGES_URL % (mirror, suite, packages_path), get_url) # type: ignore
+  pkg_map, srcpkg_map = await cache.get(
+    (cache, APT_PACKAGES_URL % (mirror, suite, packages_path)), parse_packages) # type: ignore
 
-  pkg_found = False
-  for line in apt_packages.split("\n"):
-    if pkg and line == "Package: " + pkg:
-      pkg_found = True
-    if srcpkg and line == "Source: " + srcpkg:
-      pkg_found = True
-    if pkg_found and line.startswith("Version: "):
-      version = line[9:]
-      if strip_release:
-        version = version.split("-")[0]
-      return version
+  if pkg and pkg in pkg_map:
+    version = pkg_map[pkg]
+  elif srcpkg and srcpkg in srcpkg_map:
+    version = srcpkg_map[srcpkg]
+  else:
+    raise GetVersionError('package not found in APT repository')
 
-  raise GetVersionError('package not found in APT repository')
+  if strip_release:
+    version = version.split("-")[0]
+  return version