import csv import gzip import multiprocessing as MP import os import pathlib import queue import sys from sys import exit from datetime import datetime, timedelta from urllib.request import Request, urlopen import _ed_lrr import ed_lrr_gui import requests as RQ from ed_lrr_gui import Preprocessor, Router, cfg from ed_lrr_gui.gui.ed_lrr import Ui_ED_LRR from ed_lrr_gui.gui.widget_route import Ui_diag_route from PyQt5.QtCore import QObject, Qt, QThread, QTimer, pyqtSignal from PyQt5.QtGui import QColor, QPalette, QIcon from PyQt5.QtWidgets import ( QAction, QApplication, QFileDialog, QMainWindow, QMessageBox, QProgressDialog, QTreeWidgetItem, QLabel, QDialog, ) class ProgressDialog(QProgressDialog): def __init__(self, *args, **kwargs): super().__init__(*args, **kwargs) self.setWindowModality(Qt.WindowModal) 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 Job(QObject): progress = pyqtSignal("PyQt_PyObject") def __init__(self, app, main_window, cls, *args, **kwargs): super().__init__() self.job = cls(*args, **kwargs) self.timer = QTimer(app) self.app = app self.main_window = main_window self.timer.timeout.connect(self.interval) self.timer.start(100) self.last_val = None self.progress_dialog = None self.state = {} self.setup_progress(self.handle_progess) def setup_progress(self, handle_progess): self.progress.connect(lambda *args, **kwargs: handle_progess(*args, **kwargs)) def handle_progess(self, *args, **kwargs): raise NotImplementedError def start(self): if self.progress_dialog is None: self.progress.connect( lambda *args, **kwargs: print("PROGRESS:", *args, **kwargs) ) self.started = datetime.today() return self.job.start() def cancel(self): self.job.terminate() self.job = None def __bool__(self): return not self.done @property def done(self): if self.job: return (self.job.is_alive() == False) and (self.job.queue.empty()) return True def interval(self): while self: try: res = self.job.queue.get(True, 0.1) except queue.Empty: return if res == self.last_val: continue self.state.update(res) self.progress.emit(self.state) self.last_val = res class PreprocessJob(Job): def __init__(self, app, main_window, *args, **kwargs): super().__init__(app, main_window, Preprocessor, *args, **kwargs) self.progress_dialog = ProgressDialog("", "Cancel", 0, 0, self.main_window) self.progress_dialog.setAutoClose(False) self.progress_dialog.canceled.connect(self.cancel) self.progress_dialog.show() self.start() def handle_progess(self, state): if not self: self.progress_dialog.close() return msg = "Processed: {}/{}".format( sizeof_fmt(state["status"]["done"]), sizeof_fmt(state["status"]["total"]) ) state["status"]["prc_done"] = ( state["status"]["done"] / state["status"]["total"] ) * 100 title = "[{prc_done:.2f}%] Processing {file}".format(**state["status"]) self.progress_dialog.setMinimum(0) self.progress_dialog.setMaximum(100 * 100) self.progress_dialog.setWindowTitle(title) self.progress_dialog.setLabelText(msg) self.progress_dialog.setValue(int(state["status"]["prc_done"] * 100)) class RouterJob(Job): def __init__(self, app, main_window, *args, **kwargs): super().__init__(app, main_window, Router, *args, **kwargs) self.progress_dialog = ProgressDialog("", "Cancel", 0, 0, self.main_window) self.progress_dialog.setAutoClose(False) self.progress_dialog.setLabelText("Loading (this will take a bit) ...") self.progress_dialog.setWindowTitle("Loading...") self.progress_dialog.canceled.connect(self.cancel) self.progress_dialog.show() self.start() def handle_progess(self, state): if state.get("return") is not None: self.progress_dialog.close() route_win = WRoute(self.main_window, state["return"]) return msg = "Depth: {depth}\nBody: {body}\nQueued: {queue_size}\nDistance remaining: {d_rem:.2f} Ly".format( **state["status"] ) title = "[{prc_done:.2f}%] Plotting route from [{from}] to [{to}]".format( **state["status"] ) self.progress_dialog.setMinimum(0) self.progress_dialog.setMaximum(100 * 100) self.progress_dialog.setWindowTitle(title) self.progress_dialog.setLabelText(msg) self.progress_dialog.setValue(int(state["status"]["prc_done"] * 100)) 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): dl_jobs = [ (self.systems_url, self.systems_file), (self.bodies_url, self.bodies_file), ] for url, dest in dl_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.setup_styles() def set_style(self, style): print("LOAD:", style) self.setPalette(self.styles[style]) def setup_styles(self): self.styles = {} styles = { "Dark": { "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, } } for style, colors in styles.items(): 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.styles[style] = palette self.styles["Light"] = self.style().standardPalette() class WRoute(Ui_diag_route): def __init__(self, main_window, hops): super().__init__() dialog = QDialog(main_window) self.setupUi(dialog) for n, item in enumerate(hops): if item["body"].startswith(item["system"]): item["body"] = item["body"].replace(item["system"], "").strip() item = QTreeWidgetItem( self.lst_route, [ str(n + 1), item["system"], "{body} ({star_type})".format(**item), str(item["distance"]), None, ], ) item.setFlags(item.flags() & ~Qt.ItemIsDropEnabled) dialog.exec_() class ED_LRR(Ui_ED_LRR): dl_thread = None diag_prog = None dl_started = None system_found = pyqtSignal("PyQt_PyObject") def __init__(self): super().__init__() self.current_job = None def get_open_file(self, filetypes, callback=None): fileName, _ = QFileDialog.getOpenFileName( self.main_window, "Open file", cfg["folders.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", cfg["folders.data_dir"], filetypes, options=QFileDialog.DontUseNativeDialog, ) if callback: return callback(fileName) return fileName def preprocess(self): if self.current_job: # ERROR return bodies_json = self.inp_bodies_pp.currentText() systems_json = self.inp_systems_pp.currentText() output_file = self.inp_out_pp.currentText() self.current_job = PreprocessJob( self.app, self.main_window, systems_json, bodies_json, output_file ) def set_sys_lst(self, path): if path not in cfg["history.stars_csv_path"]: cfg["history.stars_csv_path"].append(path) self.update_dropdowns() def set_bodies_file(self, path): if path not in cfg["history.bodies_path"]: cfg["history.bodies_path"].append(path) self.update_dropdowns() def set_systems_file(self, path): if path not in cfg["history.systems_path"]: cfg["history.systems_path"].append(path) self.update_dropdowns() def update_dropdowns(self): self.inp_systems_pp.clear() self.inp_systems_dest_dl.clear() for path in cfg["history.systems_path"][:]: self.inp_systems_pp.addItem(path) self.inp_systems_pp.setCurrentText(path) self.inp_systems_dest_dl.addItem(path) self.inp_systems_dest_dl.setCurrentText(path) self.inp_bodies_pp.clear() self.inp_bodies_dest_dl.clear() for path in cfg["history.bodies_path"][:]: self.inp_bodies_pp.addItem(path) self.inp_bodies_pp.setCurrentText(path) self.inp_bodies_dest_dl.addItem(path) self.inp_bodies_dest_dl.setCurrentText(path) self.inp_sys_lst.clear() self.inp_out_pp.clear() for path in cfg["history.stars_csv_path"]: self.inp_sys_lst.addItem(path) self.inp_sys_lst.setCurrentText(path) self.inp_out_pp.addItem(path) self.inp_out_pp.setCurrentText(path) 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("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.chk_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") def set_greedyness(self, value): self.lbl_greedyness.setText("Greedyness Factor ({:.0%})".format(value / 100)) @property def systems(self): ret = [] for n in range(self.lst_sys.topLevelItemCount()): ret.append(self.sys_to_dict(n)) return ret 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()) ] ret = dict(zip(header, system)) ret["id"] = getattr(self.lst_sys.topLevelItem(n), "__id__", None) ret.pop(None, None) return ret def error(self, msg): QMessageBox.critical(self.main_window, "ED_LRR Error", msg) def get_sys_list(self): if not self.inp_sys_lst.currentText(): self.error("System list is required!") return path = pathlib.Path(self.inp_sys_lst.currentText()) if not path.exists(): self.error("System list does not exist, run download and preprocess first!") return return path def compute_route(self): self.bar_status.showMessage("Resolving systems...") num_resolved = self.resolve_systems() if num_resolved: if ( QMessageBox.question( self.main_window, "ED_LRR", "Resolved {} system(s), are the names correct?".format( num_resolved ), ) == QMessageBox.No ): return self.bar_status.clearMessage() print(self.systems) systems = [str(s["id"]) for s in self.systems] jump_range = self.sb_range.value() if jump_range == 0: self.error( "Your jump range is set to zero, you're not going to get anywhere that way" ) return mode = self.cmb_mode.currentText() primary = self.chk_primary.isChecked() keep_first = self.chk_permute_keep_first.isChecked() keep_last = self.chk_permute_keep_last.isChecked() permute = self.chk_permute.isChecked() greedyness = ( self.sld_greedyness.value() / 100 if self.sld_greedyness.isEnabled() else None ) path = self.get_sys_list() if path is None: return precomp = None path = str(path) mode = { "Breadth-First Search": "bfs", "A*-Search": "astar", "Greedy-Search": "greedy", }[mode] print( systems, jump_range, mode, primary, permute, (keep_first, keep_last), greedyness, path, precomp, ) if not self.current_job: self.bar_status.showMessage("Computing Route...") self.current_job = RouterJob( self.app, self.main_window, systems, jump_range, None, mode, primary, permute, keep_first, keep_last, greedyness, precomp, path, ) else: self.error("there is already a job running!") def find_sys_by_names(self, names): t_s = datetime.today() if not self.get_sys_list(): return None # TODO: start thread/subprocess ret = _ed_lrr.find_sys(names, self.inp_sys_lst.currentText()) print("Took:", datetime.today() - t_s) return ret def resolve_systems(self): # TODO: show spinner names = [] nums = [] for n in range(self.lst_sys.topLevelItemCount()): sys_id = getattr(self.lst_sys.topLevelItem(n), "__id__", None) if sys_id is not None: continue names.append(self.sys_to_dict(n)["Name"]) nums.append(n) if not names: return 0 systems = self.find_sys_by_names(names) if systems is None: return for i, name in zip(nums, names): _, system = systems[name] self.lst_sys.topLevelItem(i).setData(0, 0, system["system"]) self.lst_sys.topLevelItem(i).setData(1, 0, system["star_type"]) self.lst_sys.topLevelItem(i).__id__ = system["id"] return len(names) # diff, item = self.find_sys_by_name(name) # print("Found", (diff, item)) def add_system(self): name = self.inp_sys.text() item = QTreeWidgetItem(self.lst_sys, [name, None]) item.resolved = False 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_dl_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_dl_progress) self.dl_thread.start() print(".") def update_permute_chk(self, state): self.chk_permute_keep_first.setEnabled(state) self.chk_permute_keep_last.setEnabled(state) self.lbl_keep.setEnabled(state) def setup_signals(self): self.btn_download.clicked.connect(self.run_download) 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.compute_route) self.btn_add.clicked.connect(self.add_system) self.btn_rm.clicked.connect(self.remove_system) self.btn_preprocess.clicked.connect(self.preprocess) self.chk_permute.stateChanged.connect(self.update_permute_chk) 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): for key in [ "history.stars_csv_path", "history.bodies_path", "history.systems_path", ]: for path in cfg[key][:]: if not os.path.isfile(path): cfg[key].remove(path) cfg.write() print("BYEEEEEE!") def setup_styles(self, win, app): for name in app.styles: action = QAction(app) action.setObjectName("action_load_style_" + name) action.setText(name) action.triggered.connect(lambda _, name=name: app.set_style(name)) self.menuStyle.addAction(action) 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() self.update_permute_chk(self.chk_permute.isChecked()) self.setup_styles(MainWindow, app) def main(): MP.freeze_support() app = App() app.setWindowIcon(QIcon(r"D:\devel\rust\ed_lrr_gui\icon\icon.ico")) MainWindow = QMainWindow() MainWindow.setWindowIcon(QIcon(r"D:\devel\rust\ed_lrr_gui\icon\icon.ico")) ui = ED_LRR() ui.setupUi(MainWindow, app) MainWindow.show() ret = app.exec_() ui.handle_close() exit(ret) if __name__ == "__main__": main()