ED_LRR/ed_lrr_gui/__main__.py

import sys
import os
import requests as RQ
from datetime import datetime, timedelta
from PyQt5.QtCore import QThread, pyqtSignal, Qt
from PyQt5.QtWidgets import (
    QMainWindow,
    QApplication,
    QFileDialog,
    QProgressDialog,
    QTreeWidgetItem,
)
from urllib.request import Request, urlopen
import gzip
from PyQt5.QtGui import QPalette, QColor
import ed_lrr_gui
import ed_lrr_gui.config as cfg
from ed_lrr_gui.gui.ed_lrr import Ui_ED_LRR
import multiprocessing as MP


def sizeof_fmt(num, suffix="B"):
    for unit in ["", "Ki", "Mi", "Gi", "Ti", "Pi", "Ei", "Zi"]:
        if abs(num) < 1024.0:
            return "{:.02f}{}{}".format(num, unit, suffix)
        num /= 1024.0
    return "{:.02f}{}{}".format(num, "Yi", suffix)


def t_round(dt):
    return dt - dt % timedelta(seconds=1)


class ProgressDialog(QProgressDialog):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.setWindowModality(Qt.WindowModal)


class DownloadThread(QThread):
    progress = pyqtSignal("PyQt_PyObject")

    def __init__(self, systems_url, systems_file, bodies_url, bodies_file):
        super().__init__()
        self.systems_url = systems_url
        self.systems_file = systems_file
        self.bodies_url = bodies_url
        self.bodies_file = bodies_file
        self.running = True

    def __del__(self):
        self.wait()

    def stop(self):
        self.running = False

    def run(self):
        jobs = [
            (self.systems_url, self.systems_file),
            (self.bodies_url, self.bodies_file),
        ]
        for url, dest in jobs:
            outfile = url.split("/")[-1]
            size = RQ.head(url, headers={"Accept-Encoding": "None"})
            size.raise_for_status()
            size = int(size.headers.get("Content-Length", 0))
            with open(dest, "wb") as of:
                resp = RQ.get(url, stream=True)
                for chunk in resp.iter_content(1024 * 1024):
                    of.write(chunk)
                    self.progress.emit(
                        {"done": of.tell(), "size": size, "outfile": outfile}
                    )
                    if not self.running:
                        return


class App(QApplication):
    def __init__(self):
        super().__init__(sys.argv)
        self.setStyle("Fusion")
        self.set_pallete()

    def set_pallete(self):
        colors = {
            "Window": QColor(53, 53, 53),
            "WindowText": Qt.white,
            "Base": QColor(15, 15, 15),
            "AlternateBase": QColor(53, 53, 53),
            "ToolTipBase": Qt.white,
            "ToolTipText": Qt.white,
            "Text": Qt.white,
            "Button": QColor(53, 53, 53),
            "ButtonText": Qt.white,
            "BrightText": Qt.red,
            "Highlight": QColor(255, 128, 0),
            "HighlightedText": Qt.black,
        }
        palette = QPalette()
        for entry, color in colors.items():
            palette.setColor(getattr(QPalette, entry), color)
            if color == Qt.darkGray:
                palette.setColor(
                    QPalette.Disabled, getattr(QPalette, entry), QColor(53, 53, 53)
                )
            else:
                palette.setColor(
                    QPalette.Disabled, getattr(QPalette, entry), Qt.darkGray
                )
        self.setPalette(palette)


class ED_LRR(Ui_ED_LRR):
    dl_thread = None
    diag_prog = None
    dl_started = None

    def __init__(self):
        super().__init__()
        self.config = cfg.load()

    def get_open_file(self, filetypes, callback=None):
        fileName, _ = QFileDialog.getOpenFileName(
            self.main_window,
            "Open file",
            str(cfg.data_dir),
            filetypes,
            options=QFileDialog.DontUseNativeDialog,
        )
        if callback:
            return callback(fileName)
        return fileName

    def get_save_file(self, filetypes, callback=None):
        fileName, _ = QFileDialog.getSaveFileName(
            self.main_window,
            "Save file",
            str(cfg.data_dir),
            filetypes,
            options=QFileDialog.DontUseNativeDialog,
        )
        if callback:
            return callback(fileName)
        return fileName

    def set_sys_lst(self, path):
        if path not in self.config.history_out_path:
            self.config.history_out_path.append(path)
            self.inp_sys_lst.addItem(path)
            self.inp_out_pp.addItem(path)

    def set_bodies_file(self, path):
        if path not in self.config.history_bodies_path:
            self.config.history_bodies_path.append(path)
            self.inp_bodies_pp.addItem(path)

    def set_systems_file(self, path):
        if path not in self.config.history_systems_path:
            self.config.history_systems_path.append(path)
            self.inp_systems_pp.addItem(path)

    def update_dropdowns(self):
        return

    def log(self, *args):
        t = datetime.today()
        msg_t = "[{}] {}".format(t, str.format(*args))
        self.txt_log.append(msg_t)

    def set_comp_mode(self, _):
        if self.rd_comp.isChecked():
            comp_mode = "Compute Route"
            self.btn_add.setText("Search+Add")
        if self.rd_precomp.isChecked():
            comp_mode = "Precompute Graph"
            self.btn_add.setText("Select")
        self.log("COMP_MODE", comp_mode)
        self.lst_sys.setEnabled(self.rd_comp.isChecked())
        self.btn_rm.setEnabled(self.rd_comp.isChecked())
        self.cmb_mode.setEnabled(self.rd_comp.isChecked())
        self.btn_permute.setEnabled(self.rd_comp.isChecked())
        self.lbl_keep.setEnabled(self.rd_comp.isChecked())
        self.lbl_mode.setEnabled(self.rd_comp.isChecked())
        self.chk_permute_keep_first.setEnabled(self.rd_comp.isChecked())
        self.chk_permute_keep_last.setEnabled(self.rd_comp.isChecked())
        self.set_route_mode(self.rd_precomp.isChecked() or None)

    def set_route_mode(self, mode=None):
        if mode == None:
            mode = self.cmb_mode.currentText()
        self.lbl_greedyness.setEnabled(mode == "A*-Search")
        self.sld_greedyness.setEnabled(mode == "A*-Search")
        self.log("ROUTE_MODE", mode)

    def set_greedyness(self, value):
        self.lbl_greedyness.setText("Greedyness Factor ({:.0%})".format(value / 100))

    def sys_to_dict(self, n):
        header = [
            self.lst_sys.headerItem().data(c, 0)
            for c in range(self.lst_sys.headerItem().columnCount())
        ]
        system = [
            self.lst_sys.topLevelItem(n).data(c, 0)
            for c in range(self.lst_sys.topLevelItem(n).columnCount())
        ]
        return dict(zip(header, system))

    def run(self):
        settings = {}
        settings["permute"] = (
            self.chk_permute_keep_first.checkState(),
            self.chk_permute_keep_last.checkState(),
        )
        settings["range"] = self.sb_range.value()
        settings["systems"] = [
            self.sys_to_dict(i) for i in range(self.lst_sys.topLevelItemCount())
        ]
        print(settings)
        progress = ProgressDialog(
            "(Not actually computing)", "Cancel", 0, 0, self.main_window
        )
        progress.setWindowTitle("Computing Route")
        progress.setFixedSize(400, 100)
        progress.setRange(0, 0)
        progress.show()

    def add_system(self):
        n = self.lst_sys.topLevelItemCount() + 1
        name = self.inp_sys.text()
        item = QTreeWidgetItem(
            self.lst_sys, [str(n) + ". Name: " + name, "Type: " + name[::-1]]
        )
        item.sys_id = "SYS_ID_HERE"
        item.setFlags(item.flags() & ~Qt.ItemIsDropEnabled)

    def remove_system(self):
        root = self.lst_sys.invisibleRootItem()
        for item in self.lst_sys.selectedItems():
            root.removeChild(item)

    def dl_canceled(self):
        if self.dl_thread:
            print("Cancel!")
            try:
                self.dl_thread.progress.disconnect()
            except TypeError:
                pass
            self.dl_thread.stop()
            self.dl_thread.wait()
            self.diag_prog.close()
            self.dl_thread = None
            self.diag_prog = None
            self.dl_started = None

    def handle_progress(self, args):
        filename = os.path.split(args["outfile"])[-1]
        if self.diag_prog is None:
            self.diag_prog = ProgressDialog("", "Cancel", 0, 1000, self.main_window)
            if self.dl_thread:
                self.diag_prog.canceled.connect(self.dl_canceled)
            self.diag_prog.show()
        t_elapsed = datetime.today() - self.dl_started
        rate = args["done"] / t_elapsed.total_seconds()
        remaining = (args["size"] - args["done"]) / rate
        rate = round(rate, 2)
        # print(rate, remaining)
        try:
            t_rem = timedelta(seconds=remaining)
        except OverflowError:
            t_rem = "-"
        msg = "Downloading {} [{}/{}] ({}/s)\n[{}/{}]".format(
            filename,
            sizeof_fmt(args["done"]),
            sizeof_fmt(args["size"]),
            sizeof_fmt(rate),
            t_round(t_elapsed),
            t_round(t_rem),
        )
        self.diag_prog.setLabelText(msg)
        self.diag_prog.setWindowTitle("Downloading EDSM Dumps")
        self.diag_prog.setValue((args["done"] * 1000) // args["size"])

    def run_download(self):
        if self.dl_thread:
            return
        self.dl_started = datetime.today()
        self.dl_thread = DownloadThread(
            self.inp_systems_dl.currentText(),
            self.inp_systems_dest_dl.currentText(),
            self.inp_bodies_dl.currentText(),
            self.inp_bodies_dest_dl.currentText(),
        )
        self.dl_thread.progress.connect(self.handle_progress)
        self.dl_thread.start()
        print(".")

    def setup_signals(self):
        self.btn_download.clicked.connect(self.run_download)
        self.inp_systems_dest_dl.setCurrentText(r"D:\devel\rust\ed_lrr_gui\DL\s.json")
        self.inp_bodies_dest_dl.setCurrentText(r"D:\devel\rust\ed_lrr_gui\DL\b.json")
        self.set_greedyness(self.sld_greedyness.value())
        self.cmb_mode.currentTextChanged.connect(self.set_route_mode)
        self.rd_comp.toggled.connect(self.set_comp_mode)
        self.rd_precomp.toggled.connect(self.set_comp_mode)
        self.sld_greedyness.valueChanged.connect(self.set_greedyness)
        self.btn_go.clicked.connect(self.run)
        self.btn_add.clicked.connect(self.add_system)
        self.btn_rm.clicked.connect(self.remove_system)
        self.btn_out_browse_pp.clicked.connect(
            lambda: self.get_save_file("CSV File (*.csv)", self.set_sys_lst)
        )
        self.btn_sys_lst_browse.clicked.connect(
            lambda: self.get_open_file("CSV File (*.csv)", self.set_sys_lst)
        )

        self.btn_bodies_browse_pp.clicked.connect(
            lambda: self.get_open_file("JSON File (*.json)", self.set_bodies_file)
        )
        self.btn_bodies_dest_browse_dl.clicked.connect(
            lambda: self.get_save_file("JSON File (*.json)", self.set_bodies_file)
        )
        self.btn_systems_browse_pp.clicked.connect(
            lambda: self.get_open_file("JSON File (*.json)", self.set_systems_file)
        )
        self.btn_systems_dest_browse_dl.clicked.connect(
            lambda: self.get_save_file("JSON File (*.json)", self.set_systems_file)
        )

    def handle_close(self):
        cfg.write(self.config)
        print("BYEEEEEE!")

    def setupUi(self, MainWindow, app):
        super().setupUi(MainWindow)
        self.update_dropdowns()
        self.main_window = MainWindow
        self.app = app
        self.setup_signals()
        self.lst_sys.setHeaderLabels(["Name", "Type"])
        self.set_route_mode()


def main():
    app = App()
    MainWindow = QMainWindow()
    ui = ED_LRR()
    ui.setupUi(MainWindow, app)
    MainWindow.show()
    ret = app.exec_()
    ui.handle_close()
    exit(ret)


if __name__ == "__main__":
    MP.freeze_support()
    main()