๐Ÿ“ฆ anlutro / psm

๐Ÿ“„ psm.sh ยท 204 lines
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204#!/bin/bash
set -eu

[ -z "${PSM_VENV_DIR-}" ] && PSM_VENV_DIR=~/.local/share/psm
[ -z "${PSM_BIN_DIR-}" ] && PSM_BIN_DIR=~/.local/bin
[ -z "${PSM_PYTHON-}" ] && PSM_PYTHON=$(
    find $HOME/.local/bin /usr/local/bin /usr/bin -regex '.*/python[3-9][0-9.]*' -printf '%f\n' \
    | sort -V | tail -1
)
PSM_PYTHON_VER=$($PSM_PYTHON --version 2>&1 | cut -d' ' -f2)

export PIP_DISABLE_PIP_VERSION_CHECK=1

_get_pkg_name() {
    if [ -d "$1" ] && [ -e "$1/setup.py" ]; then
        $PSM_PYTHON "$1/setup.py" --name
    else
        # no longer available in stdlib / by default
        # $PSM_PYTHON -c "from packaging.requirements import Requirement; print(Requirement('$1').name)"
        echo "$1" | grep -oP '^[a-zA-Z0-9-_]+'
    fi
}

_psm_list() {
    {
        for venv in $PSM_VENV_DIR/*/; do
            name=$(basename $venv)
            if [ -e $venv/bin/python ]; then
                $venv/bin/pip list "$@" | grep -P "^$name\s+" || true
            else
                echo >&2 "WARNING: venv for $name is broken! to fix: psm reinstall $name"
            fi
        done
    } | column -t
}

_psm_list_all_scripts() {
    if [ $# -gt 0 ]; then
        pkgs="$*"
    else
        pkgs=$(find $PSM_VENV_DIR -mindepth 1 -maxdepth 1 -type d -printf '%f\n' | sort)
    fi
    for pkg in $pkgs; do
        scripts=$(_psm_list_scripts $pkg)
        if [ -n "$scripts" ]; then
            echo $pkg: $scripts
        fi
    done
}

_psm_list_scripts() {
    if [ -z "$1" ]; then
        echo "missing argument!" >&2
        return 1
    fi
    venv=$PSM_VENV_DIR/$1
    if [ ! -d "$venv" ]; then
        echo >&2 "venv does not exist: $venv"
        return 1
    fi
    if [ ! -e "$venv/bin/python" ]; then
        echo >&2 "bin/python is broken for venv: $venv"
        return 1
    fi
    $venv/bin/python -c "
from os.path import abspath, basename, join
import importlib.metadata
dist = importlib.metadata.distribution('$1')
for script in dist.entry_points.select(group='console_scripts'):
    print(script.name)
installed = dist._read_files_egginfo_installed()
if installed:
    for line in installed:
        path = abspath(join(dist.locate_file(''), line.split(',')[0]))
        if path.startswith('$venv/bin/'):
            print(basename(path))
records = dist._read_files_distinfo()
if records:
    records = [s.split(',')[0] for s in records]
    bin_paths = [s for s in records if '/bin/' in s and not s.endswith('.pyc')]
    for bin_path in bin_paths:
        print(bin_path.split('/')[-1])
" | sort | uniq
}

_install_cleanup() {
    rm -rf ${PSM_VENV_DIR:?}/$pkg_name
}

_psm_install() {
    for pkg in "$@"; do
        pkg_name=$(_get_pkg_name "$pkg")
        if [ -d $PSM_VENV_DIR/$pkg_name ]; then
            echo "$pkg_name already exists, not doing anything."
            echo "Hint: did you mean psm upgrade or psm reinstall?"
        else
            trap _install_cleanup EXIT
            echo "Creating virtual environment for $pkg_name with $PSM_PYTHON ..."
            $PSM_PYTHON -m venv $PSM_VENV_DIR/$pkg_name || exit 1
            _psm_upgrade "$pkg"
            trap - EXIT
        fi
    done
}

_psm_reinstall() {
    for pkg in "$@"; do
        _psm_uninstall "$pkg"
        _psm_install "$pkg"
    done
}

_upgade_cleanup() {
    echo "PSM script aborted in the middle of upgrade!"
    echo "$venv may be in an inconsistent state."
    echo "If broken, run:  psm reinstall '$pkg'"
}

_psm_upgrade() {
    for pkg in "$@"; do
        pkg_name=$(_get_pkg_name "$pkg")
        venv=$PSM_VENV_DIR/$pkg_name

        # determine if package is editable
        if [ -d "$pkg" ] && [ -e "$pkg/setup.py" ]; then
            editable_dir=$(readlink -f "$pkg")
        else
            editable_dir=$($venv/bin/pip list --editable | awk "\$1 == \"$pkg_name\" { print \$3 }")
        fi

        trap _upgade_cleanup EXIT

        # check if venv can be upgraded with new python
        venv_pyver=$($venv/bin/python --version 2>&1 | cut -d' ' -f2)
        if [ $venv_pyver != $PSM_PYTHON_VER ]; then
            echo "Recreating venv with new python ($PSM_PYTHON) for $pkg_name ..."
            $PSM_PYTHON -m venv --clear $venv
        fi

        echo "Installing/upgrading pip and setuptools for $pkg_name ..."
        $venv/bin/pip install -q -U pip setuptools

        echo "Installing package: $pkg_name ..."
        if [ -n "$editable_dir" ]; then
            pip_install_args=(-e "$editable_dir")
        else
            pip_install_args=("$pkg")
        fi
        $venv/bin/pip install -q -U "${pip_install_args[@]}"

        echo "Creating script symlinks for $pkg_name ..."
        _psm_list_scripts $pkg_name | xargs -r -I% ln -sf $venv/bin/% $PSM_BIN_DIR/

        trap - EXIT
    done
}

_psm_upgrade_all() {
    pkgs=$(
        find $PSM_VENV_DIR -mindepth 1 -maxdepth 1 -type d -print0 \
        | xargs -r -0 -n1 basename
    )
    _psm_upgrade $pkgs
}

_psm_uninstall() {
    for pkg in "$@"; do
        echo "Uninstalling package: $pkg ..."
        _psm_list_scripts $pkg | xargs -r -I% rm -f $PSM_BIN_DIR/%
        rm -rf "${PSM_VENV_DIR:?}/$pkg"
    done
}

_psm_help() {
    echo "psm - python script manager"
    echo
    echo "psm install pkg [pkg ...]"
    echo "psm uninstall pkg [pkg ...]"
    echo "psm upgrade pkg [pkg ...]"
    echo "psm upgrade-all"
}

psm() {
    if [ $# -gt 0 ]; then
        arg="$1" && shift
        cmd="$(echo "$arg" | tr -s - _)"
    fi

    if [ -z "${cmd-}" ] || [ "$arg" = '-h' ] || [ "$arg" = '--help' ]; then
        func=_psm_help
    else
        func="_psm_$cmd"
    fi

    if ! type $func >/dev/null 2>&1; then
        echo "Unknown command $arg"
        exit 1
    fi

    $func "$@"
}

psm "$@"