.pipi.py

#!/usr/bin/env python3
import os
import subprocess
import sys
import typing

if typing.TYPE_CHECKING:
    from typing import Optional  # NOQA:F401

"""
Custom install shortcut as i for pip

Usage::
    $ python .pipi.py

"""

IS_WINDOWS = sys.platform == "win32"
TEXT = """
    if sys.argv[1:]:
        is_installing = sys.argv[1] == 'install'
        if not is_installing and sys.argv[1] == 'i':
            sys.argv[1] = 'install'
            is_installing = True
        if is_installing and sys.argv[2:]:
            a2 = sys.argv[2]
            s = 'git+https://'
            if a2.startswith(s):
                sys.argv[2] = a2.replace(s, 'git+ssh://git@')
"""


def run_and_echo(cmd):
    # type: (str) -> int
    print("--> " + cmd)
    sys.stdout.flush()
    return os.system(cmd)


def capture_output(cmd):
    # type: (str) -> str
    try:
        r = subprocess.run(cmd, shell=True, capture_output=True)
    except (TypeError, AttributeError):  # For python<=3.6
        with os.popen(cmd) as p:
            return p.read().strip()
    else:
        return r.stdout.decode().strip()


def patch_it(filename, tip="pip i package-name"):
    # type: (str, str) -> Optional[int]
    if IS_WINDOWS:
        # TODO: change pip main
        return 0
    with open(filename) as f:
        text = f.read()
    if "install" not in text:
        ss = text.strip().splitlines()
        new_text = "\n".join(ss[:-1] + [TEXT] + ss[-1:])
        try:
            with open(filename, "w") as f:
                f.write(new_text)
        except IOError:
            print("Failed to write lines to {}:\n{}".format(filename, TEXT))
            return 1
        print("i command was configured!\n\nUsage::\n\n    {}\n".format(tip))


def main():
    # type: () -> Optional[int]
    pip_file = os.path.join(os.path.dirname(sys.executable), "pip")
    if not os.path.exists(pip_file):
        pip_file = capture_output("which pip")
    args = sys.argv[1:]
    packages = " ".join(args)
    if patch_it(pip_file):
        return 1
    elif not args:
        if IS_WINDOWS:
            run_and_echo("pipx inject poetry poetry-plugin-i")
        else:
            poetry_file = capture_output("which poetry")
            if poetry_file and os.path.exists(poetry_file):
                patch_it(poetry_file, "poetry i")
        return run_and_echo("pip install")
    if args:
        cmd = "pip i " + " ".join([repr(i) for i in args])
        run_and_echo(cmd)
        print("Next time you can use this instead:\n\n  pip i {}\n".format(packages))
        if "pip" in args:
            return patch_it(pip_file)


if __name__ == "__main__":
    sys.exit(main())