Skip to content

Instantly share code, notes, and snippets.

@bulletinmybeard
Last active March 12, 2024 19:57
Show Gist options
  • Save bulletinmybeard/5a1ed685af72709791c576e02206603a to your computer and use it in GitHub Desktop.
Save bulletinmybeard/5a1ed685af72709791c576e02206603a to your computer and use it in GitHub Desktop.
Zsh function to extend the `poetry` command

Extend the poetry command...

function poetry() {
    function ensure_virtual_environment() {
        local venv_path=".venv"

        if [[ ! -d $venv_path ]]; then
            echo -e "Virtual environment not found."
            python3 -m venv $venv_path            
            echo -e "Virtual environment created at \033[32m$venv_path\033[0m."
        else
            echo -e "Using existing virtual environment at \033[32m$venv_path\033[0m."
        fi

        source $venv_path/bin/activate
    }

    function handle_poetry_output_and_move_configs() {
        local output=$(poetry "$@" 2>&1)

        if echo "$output" | grep -q 'Configuration file exists at'; then
            local first_path=$(echo "$output" | grep -oE '/[^,]*pypoetry' | head -1)
            local second_path=$(echo "$output" | grep -oE '/[^,]*pypoetry' | tail -1)

            if [[ -n "$first_path" && -n "$second_path" ]]; then
                find "$first_path" -name '*.toml' -print0 | xargs -0 -I {} mv {} "$second_path/" 2>/dev/null
                if [[ $? -eq 0 ]]; then
                    echo "TOML configuration files moved from $first_path to $second_path"
                    return 0
                else
                    echo "Failed to move TOML configuration files."
                    return 1
                fi
            else
                echo "Could not extract configuration paths from the output."
                return 1
            fi
        fi
    }

    function switch_global_poetry_version() {
        if [[ -z $1 ]]; then
            echo "No version provided. Usage: poetry switch-global-version <version>"
            return 1
        fi

        local target_version="$1"
        local local_version=$(poetry -V 2>&1 | grep -o 'version [0-9]*\.[0-9]*\.[0-9]*' | awk '{print $2}')

        if [[ $target_version == $local_version ]]; then
            echo "The target version $target_version is already the current version."
            return 1
        fi

        if curl -sSL https://install.python-poetry.org | python - --version $target_version; then
            local new_local_version=$(handle_poetry_output local-poetry-version)
            if [[ $target_version == $new_local_version ]]; then
                echo "Poetry version successfully switched to $target_version."
                return 0
            else
                echo "Failed to switch Poetry version to $target_version."
                return 1
            fi
        else
            echo "Failed to switch Poetry version to $target_version."
            return 1
        fi
    }

    function switch_poetry_version() {
        if [[ -z $1 ]]; then
            echo -e "\033[31mNo version provided.\033[0m Usage: poetry switch-version \033[1m\033[97m<version>\033[0m"
            return 1
        fi

        local target_version="$1"
        local local_version=$(poetry -V 2>&1 | grep -o 'version [0-9]*\.[0-9]*\.[0-9]*' | awk '{print $2}')

        if [[ $target_version == $local_version ]]; then
            echo "The target version \033[1m\033[97m$target_version\033[0m is already the current version."
            return 1
        fi

        ensure_virtual_environment

        pip install "poetry==$target_version" > /dev/null

        if (( $? != 0 )); then
            echo -e "Failed to install Poetry version \033[1m\033[97m$target_version\033[0m"
            return 1
        fi

        local new_local_version=$(poetry -V 2>&1 | grep -o 'version [0-9]*\.[0-9]*\.[0-9]*' | awk '{print $2}')

        if [[ $target_version == $new_local_version ]]; then
            echo -e "Poetry version successfully switched to \033[1m\033[97m$target_version\033[0m within the virtual environment."
            return 0
        else
            echo -e "Failed to switch Poetry version to \033[1m\033[97m$target_version\033[0m."
            return 1
        fi
    }

    function remove_poetry_version() {
        if pip list 2>&1 | grep -v 'WARNING: Skipping' | grep -q 'poetry'; then
            local current_version=$(poetry -V 2>&1 | grep -o 'version [0-9]*\.[0-9]*\.[0-9]*' | awk '{print $2}')

            if [[ -z $current_version ]]; then
                echo -e "\033[31mCurrent Poetry version could not be determined.\033[0m"
                return 1
            fi

            echo "Attempting to uninstall Poetry version \033[1m\033[97m$current_version\033[0m from the virtual environment..."
            pip uninstall --yes "poetry==$current_version" > /dev/null

            if (( $? != 0 )); then
                echo -e "\033[31mFailed to uninstall Poetry version \033[1m\033[97m$current_version\033[0m.\033[0m"
                return 1
            else
                echo -e "Poetry version \033[1m\033[97m$current_version\033[0m has been successfully uninstalled from the virtual environment."
                return 0
            fi
        else
            echo -e "\033[31mPoetry is not installed in the virtual environment.\033[0m"
            return 1
        fi
    }

    if [[ $1 == "switch-global-version" ]]; then
        shift
        handle_poetry_output_and_move_configs
        switch_global_poetry_version "$@"
    elif [[ $1 == "switch-version" ]]; then
        shift
        handle_poetry_output_and_move_configs
        switch_poetry_version "$@"
    elif [[ $1 == "remove-version" ]]; then
        shift
        handle_poetry_output_and_move_configs
        remove_poetry_version
    elif [[ $1 == "--version" ]]; then
        handle_poetry_output_and_move_configs
        command poetry "$@"
    else
        command poetry "$@"
    fi
}
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment