utils.py

import re
import os
import shlex
import shutil
import urllib
import tarfile
import inquirer
import requests
import tempfile
import zstandard
import subprocess

from pathlib import Path
from functools import partial
from urllib.request import urlopen
from rich.progress import (
    BarColumn,
    DownloadColumn,
    Progress,
    TextColumn,
    TimeRemainingColumn,
    TransferSpeedColumn,
)

from config import *


class Utils:
    def __init__(self) -> None:
        self.aur_pkgs = {}
        if USE_AUR:
            self.get_all_pkgs_aur()

    def new_progress(self):
        return Progress(
            TextColumn("[bold blue]{task.fields[filename]}", justify="right"),
            BarColumn(bar_width=None),
            DownloadColumn(),
            "•",
            TransferSpeedColumn(),
            "•",
            TimeRemainingColumn(),
            "  "
        )

    def get_pkg(self, pkg_name):
        output = None
        data = requests.get(ARCH_SEARCH_URL.format(pkg=pkg_name)).json()
        res = data["results"]

        if len(res) > 0:
            output = ARCH_URL.format(
                repo=res[0]["repo"],
                arch=res[0]["arch"],
                pkg=res[0]["pkgname"]
            )
            return output

        aur_matches = {pkg: self.aur_pkgs[pkg]
                       for pkg in self.aur_pkgs if pkg.startswith(pkg_name)}

        if len(aur_matches) == 1:
            output = os.path.join(AUR_URL, list(aur_matches.values())[0])

        elif len(aur_matches) > 1:
            choices = list(aur_matches.keys())
            choices.append("None of the above")
            s = self.user_select(
                "Multiple Packages were found. Please select which one to use", choices)
            if s != "None of the above":
                output = os.path.join(AUR_URL, aur_matches[s])

        return output

    def user_select(self, que, choices):
        print(f"\n{que}")
        q = [inquirer.List('a', message=">>", choices=choices)]
        return inquirer.prompt(q)["a"]

    def user_text(self, que):
        print(f"\n{que}")
        q = [inquirer.Text('a', message=">>")]
        return inquirer.prompt(q)["a"]

    def user_confirm(self, que, y_otp="Yes", n_opt="No"):
        ans = self.user_select(que, [y_otp, n_opt])
        return ans == y_otp

    def user_path(self, que, required=False):
        while True:
            print(f"\n{que}")
            q = [inquirer.Text('a', message=">>")]
            p = inquirer.prompt(q)["a"]
            if not p:
                if required:
                    print("This field is required.")
                    continue
                p = None
                break

            if os.path.exists(p):
                break

            print("File does not exist. Please try again.\n")

        return p

    def get_all_pkgs_aur(self):
        print("Loading Chaotic AUR package list...")
        txt = requests.get(AUR_URL).text
        x = re.findall('href="(.*tar.zst)"', txt)

        for i in x:
            name = urllib.parse.unquote(
                i, encoding='utf-8', errors='replace').replace(".pkg.tar.zst", "")
            self.aur_pkgs[name] = i

    def download(self, url, dest, name):
        progress = self.new_progress()
        with progress:
            print(f"Downloading {name}...")
            task_id = progress.add_task("download", filename="  ", start=False)
            response = urlopen(url)
            progress.update(task_id, total=int(
                response.info()["Content-length"]))
            with open(dest, "wb") as dest_file:
                progress.start_task(task_id)
                for data in iter(partial(response.read, 32768), b""):
                    dest_file.write(data)
                    progress.update(task_id, advance=len(data))

    def extract_zst(self, zst_file, out_path):
        zst_file = Path(zst_file).expanduser()
        out_path = Path(out_path).expanduser().resolve()
        dctx = zstandard.ZstdDecompressor()

        with tempfile.TemporaryFile(suffix=".tar") as ofh:
            with zst_file.open("rb") as ifh:
                dctx.copy_stream(ifh, ofh)
            ofh.seek(0)
            with tarfile.open(fileobj=ofh) as z:
                z.extractall(out_path)

    def validate_desktop_file(self, file_path):
        if out := self.run_cmd(f"desktop-file-validate {file_path}", True):
            return False, out
        else:
            return True, None

    def copy_file(self, src_file, dest_dir):
        file_name = src_file.split("/")[-1]
        dest_file = os.path.join(dest_dir, file_name)
        if src_file != dest_file:
            shutil.copyfile(src_file, dest_file)

    def run_cmd(self, cmd, record_output=False):
        if record_output:
            proc = subprocess.run(shlex.split(cmd), stdout=subprocess.PIPE)
            return proc.stdout.decode("utf-8").strip()
        else:
            subprocess.run(shlex.split(cmd))

    def set_icon_desktop_file(self, desktop_file, icon_name):
        self.run_cmd(
            f"desktop-file-edit {desktop_file} --set-icon={icon_name}")

    def make_executable(self, file_path):
        self.run_cmd(f"chmod +x {file_path}")

    def rm(self, f):
        if os.path.exists(f):
            os.remove(f)

    def extract_deps(self, pkginfo_file):
        pkgs = {}
        with open(pkginfo_file, "r") as f:
            lines = f.read().splitlines()
            for line in lines:
                line = line.strip()
                if line.startswith("depend ="):
                    pkgs[line.replace("depend =", "").strip()] = None

        return pkgs

    def max_len(self, lst):
        return max(len(i) for i in lst)