From cd1cbfde308323de26bbf7d152711b0c616c7b7a Mon Sep 17 00:00:00 2001
From: lilydjwg <lilydjwg@gmail.com>
Date: Thu, 24 Sep 2020 18:27:30 +0800
Subject: [PATCH] support per source plugin configuration and support different
 registry for the npm source

---
 docs/plugin.rst         | 15 +++++++++++++++
 docs/usage.rst          |  6 ++++++
 nvchecker/__main__.py   |  1 +
 nvchecker/core.py       | 15 +++++++++++++--
 nvchecker_source/npm.py |  5 +++++
 tests/conftest.py       |  2 +-
 6 files changed, 41 insertions(+), 3 deletions(-)

diff --git a/docs/plugin.rst b/docs/plugin.rst
index 3cd500d..df0de6c 100644
--- a/docs/plugin.rst
+++ b/docs/plugin.rst
@@ -79,3 +79,18 @@ put results in :attr:`result_q <nvchecker.api.BaseWorker.result_q>`. See
 ``nvchecker_source/aur.py`` for a complete, batching example.
 
 For details about these objects, see :mod:`the API documentation <nvchecker.api>`.
+
+You can also receive a configuration section from the configuration as
+``__config__.source.SOURCE_NAME``, where ``SOURCE_None`` is what your plugin is
+called. This can be used to specify a mirror site for your plugin to use, e.g.
+the ``npm`` plugin accepts the following config::
+
+  [__config__.source.npm]
+  registry = "https://registry.npm.taobao.org"
+
+When such a configuration exists for your plugin, you need to define a function
+named ``configure`` to receive it::
+
+  def configure(config):
+    '''use the "config" dict in some way'''
+    ...
diff --git a/docs/usage.rst b/docs/usage.rst
index be6ba29..f6265fd 100644
--- a/docs/usage.rst
+++ b/docs/usage.rst
@@ -446,6 +446,12 @@ Check `NPM Registry <https://registry.npmjs.org/>`_ for updates.
 npm
   The name used on NPM Registry, e.g. ``coffee-script``.
 
+To configure which registry to query, a source plugin option is available.
+You can specify like this::
+
+  [__config__.source.npm]
+  registry = "https://registry.npm.taobao.org"
+
 Check Hackage
 ~~~~~~~~~~~~~
 ::
diff --git a/nvchecker/__main__.py b/nvchecker/__main__.py
index 76d0c59..5cb9af9 100755
--- a/nvchecker/__main__.py
+++ b/nvchecker/__main__.py
@@ -50,6 +50,7 @@ def main() -> None:
     futures = core.dispatch(
       entries, task_sem, result_q,
       keymanager, args.tries,
+      options.source_configs,
     )
   except ModuleNotFoundError as e:
     sys.exit(f'Error: {e}')
diff --git a/nvchecker/core.py b/nvchecker/core.py
index 0dec714..4634bd2 100644
--- a/nvchecker/core.py
+++ b/nvchecker/core.py
@@ -11,7 +11,7 @@ import logging
 import argparse
 from typing import (
   Tuple, NamedTuple, Optional, List, Union,
-  cast, Dict, Awaitable, Sequence,
+  cast, Dict, Awaitable, Sequence, Any,
 )
 import types
 from pathlib import Path
@@ -139,6 +139,7 @@ class Options(NamedTuple):
   max_concurrency: int
   proxy: Optional[str]
   keymanager: KeyManager
+  source_configs: Dict[str, Dict[str, Any]]
 
 class FileLoadError(Exception):
   def __init__(self, kind, exc):
@@ -159,6 +160,7 @@ def load_file(
 
   ver_files: Optional[Tuple[Path, Path]] = None
   keymanager = KeyManager(None)
+  source_configs = {}
 
   if '__config__' in config:
     c = config.pop('__config__')
@@ -184,6 +186,9 @@ def load_file(
       except OSError as e:
         raise FileLoadError('keyfile', e)
 
+    if 'source' in c:
+      source_configs = c['source']
+
     max_concurrency = c.get('max_concurrency', 20)
     proxy = c.get('proxy')
   else:
@@ -191,7 +196,9 @@ def load_file(
     proxy = None
 
   return cast(Entries, config), Options(
-    ver_files, max_concurrency, proxy, keymanager)
+    ver_files, max_concurrency, proxy, keymanager,
+    source_configs,
+  )
 
 def dispatch(
   entries: Entries,
@@ -199,6 +206,7 @@ def dispatch(
   result_q: Queue[RawResult],
   keymanager: KeyManager,
   tries: int,
+  source_configs: Dict[str, Dict[str, Any]],
 ) -> List[asyncio.Future]:
   mods: Dict[str, Tuple[types.ModuleType, List]] = {}
   ctx_tries.set(tries)
@@ -210,6 +218,9 @@ def dispatch(
       mod = import_module('nvchecker_source.' + source)
       tasks: List[Tuple[str, Entry]] = []
       mods[source] = mod, tasks
+      config = source_configs.get(source)
+      if config and getattr(mod, 'configure'):
+        mod.configure(config) # type: ignore
     else:
       tasks = mods[source][1]
     tasks.append((name, entry))
diff --git a/nvchecker_source/npm.py b/nvchecker_source/npm.py
index f2a3a12..cc4102f 100644
--- a/nvchecker_source/npm.py
+++ b/nvchecker_source/npm.py
@@ -7,6 +7,11 @@ from nvchecker.api import session
 
 NPM_URL = 'https://registry.npmjs.org/%s'
 
+def configure(config):
+  global NPM_URL
+  if url := config.get('registry'):
+    NPM_URL = f'{url.rstrip("/")}/%s'
+
 async def get_first_1k(url):
   headers = {
     "Accept": "application/vnd.npm.install-v1+json",
diff --git a/tests/conftest.py b/tests/conftest.py
index 741d989..f02c92c 100644
--- a/tests/conftest.py
+++ b/tests/conftest.py
@@ -29,7 +29,7 @@ async def run(
 
   futures = core.dispatch(
     entries, task_sem, result_q,
-    keymanager, 1,
+    keymanager, 1, {},
   )
 
   oldvers: VersData = {}