Tool for setting the index-url of pip

pip_index.py

#!/usr/bin/env python3
# coding: utf-8

"""\
Tool for setting the index-url of pip

Reference: 
- https://pip.pypa.io/en/latest/
- https://docs.python.org/3/installing/index.html
- https://docs.python.org/3/library/ensurepip.html
"""

assert __name__ == "__main__", "Can only be used as __main__ module"

__author__ = "ChenyangGao <https://chenyanggao.github.io/>"
__version__ = (0, 1)

from argparse import ArgumentParser, RawTextHelpFormatter


parser = ArgumentParser(
    description="Python pip 源配置工具", 
    formatter_class=RawTextHelpFormatter, 
)
subparsers = parser.add_subparsers(dest="command", help="Available commands")

parser_list = subparsers.add_parser(
    "list", formatter_class=RawTextHelpFormatter, 
    help="罗列所有可用的源", description="""罗列所有可用的源，格式形如
{key}: 
    index-url: {index_url!r}
    trusted-host: {trusted_host!r}
    comment: {comment!r}""")

parser_use = subparsers.add_parser(
    "use", formatter_class=RawTextHelpFormatter, 
    help="使用 key 指定要设置的源", description="使用 key 指定要设置的源")
parser_use.add_argument(
    "key", default="pypi", nargs="?", help="key，对应某个源，默认值 pypi")
parser_use.add_argument(
    "-k", "--kind", default="user", choices=("global", "user", "site"), 
    help="""类型，默认值 user
  - global  Use the system-wide configuration file only
  - user    Use the user configuration file only
  - site    Use the current environment configuration file only""",
)
parser_use.add_argument(
    "-i", "--isolated", action="store_true",
    help="Run pip in an isolated mode, ignoring environment "
         "variables and user configuration.")
parser_use.add_argument(
    "-c", "--cmd-only", dest="cmd_only", action="store_true",
    help="并不设置配置文件，只是打印一个 pip install 命令"
)

args = parser.parse_args()


from pip._internal.configuration import Configuration
from urllib.parse import urlsplit


# TIPS: 20 == logging.INFO
__import__("logging").basicConfig(level=20, format="%(message)s")

INDEXES = {
    "pypi": {
        "index-url": "https://pypi.python.org/pypi",
        "trusted-host": "pypi.python.org",
        "comment": "官方源"
    },
    "douban": {
        "index-url": "http://pypi.douban.com/simple/",
        "trusted-host": "pypi.douban.com",
        "comment": "豆瓣"
    },
    "aliyun": {
        "index-url": "http://mirrors.aliyun.com/pypi/simple/",
        "trusted-host": "mirrors.aliyun.com",
        "comment": "阿里云"
    },
    "tsinghua": {
        "index-url": "https://pypi.tuna.tsinghua.edu.cn/simple/",
        "trusted-host": "pypi.tuna.tsinghua.edu.cn",
        "comment": "清华大学"
    },
    "ustc": {
        "index-url": "https://pypi.mirrors.ustc.edu.cn/simple/",
        "trusted-host": "pypi.mirrors.ustc.edu.cn",
        "comment": "中国科学技术大学"
    },
    "hustunique": {
        "index-url": "http://pypi.hustunique.com/simple/",
        "trusted-host": "pypi.hustunique.com",
        "comment": "华中理工大学"
    },
    "sdutlinux": {
        "index-url": "http://pypi.sdutlinux.org/simple/",
        "trusted-host": "pypi.sdutlinux.org",
        "comment": "山东理工大学"
    }, 
    "tencent": {
        "index-url": "http://mirrors.cloud.tencent.com/pypi/simple/",
        "trusted-host": "mirrors.cloud.tencent.com",
        "comment": "腾讯云"
    }, 
}


def pipi_list():
    "罗列所有可用的源"
    for key, index in INDEXES.items():
        print("%s: %s" % (key,
            "".join("\n    %s: %r" % e for e in index.items())))


def pipi_use(key="pypi", kind="user", isolated=False):
    """使用 key 指定要设置的源

    :param key: key，对应某个源，默认值 pypi
    :param kind: 类型，默认值 user
        - global  Use the system-wide configuration file only
        - user    Use the user configuration file only
        - site    Use the current environment configuration file only
    :param isolated: Run pip in an isolated mode, ignoring environment 
        variables and user configuration.
    """
    # assert kind in ("global", "user", "site")
    index = INDEXES[key]
    index_url = index["index-url"]
    trusted_host = index.get("trusted-host") or urlsplit(index_url).netloc
    conf = Configuration(isolated, kind)
    conf.load()
    conf.set_value("global.index-url", index_url)
    conf.set_value("install.trusted-host", trusted_host)
    conf.save()


command = args.command
if command == "list":
    pipi_list()
elif command == "use":
    if args.cmd_only:
        from sys import executable

        index = INDEXES[args.key]
        index_url = index["index-url"]
        trusted_host = index.get("trusted-host") or urlsplit(index_url).netloc
        print(f"'{executable}' -m pip install --index-url "
              f"{index_url} --trusted-host {trusted_host} ")
    else:
        pipi_use(args.key, args.kind, args.isolated)
else:
    raise NotImplementedError(f"Command {command!r} is not implemented!")