Lgndr: Use custom wrapped decorator to wrap LegendaryCLI functions

Lgndr: Add `get_boolean_choice` to relevant args dataclasses
Lgndr: Move mock functions to `api_monkeys`

InstallDialog: Add status queue to prepare_overlay_install arguments, fixes missing download stats

rare/components/dialogs/install_dialog.py

@@ -9,7 +9,7 @@ from PyQt5.QtGui import QCloseEvent, QKeyEvent
 from PyQt5.QtWidgets import QDialog, QFileDialog, QCheckBox
 
 from rare.lgndr.api_arguments import LgndrInstallGameArgs
-from rare.lgndr.cli import LegendaryCLI
+from rare.lgndr.api_exception import LgndrException
 from rare.lgndr.core import LegendaryCore
 from legendary.models.downloading import ConditionCheckResult
 from legendary.models.game import Game
@@ -312,7 +312,7 @@ class InstallInfoWorker(QRunnable):
                 cli = LegendaryCLISingleton()
                 download = InstallDownloadModel(
                     # *self.core.prepare_download(
-                    *cli.prepare_install(LgndrInstallGameArgs(
+                    *cli.install_game(LgndrInstallGameArgs(
                         app_name=self.dl_item.options.app_name,
                         base_path=self.dl_item.options.base_path,
                         force=self.dl_item.options.force,
@@ -347,6 +347,7 @@ class InstallInfoWorker(QRunnable):
 
                 dlm, analysis, igame = self.core.prepare_overlay_install(
                     path=self.dl_item.options.base_path,
+                    status_q=self.dl_item.status_q,
                 )
 
                 download = InstallDownloadModel(
@@ -363,6 +364,8 @@ class InstallInfoWorker(QRunnable):
                 self.signals.result.emit(download)
             else:
                 self.signals.failed.emit("\n".join(str(i) for i in download.res.failures))
+        except LgndrException as ret:
+            self.signals.failed.emit(ret.message)
         except Exception as e:
             self.signals.failed.emit(str(e))
         self.signals.finished.emit()

rare/lgndr/api_arguments.py

@@ -2,6 +2,8 @@ from dataclasses import dataclass
 from multiprocessing import Queue
 from typing import Callable, List
 
+from .api_monkeys import get_boolean_choice
+
 
 @dataclass(kw_only=True)
 class LgndrCommonArgs:
@@ -21,6 +23,8 @@ class LgndrImportGameArgs:
     skip_dlcs: bool = False
     with_dlcs: bool = False
     yes: bool = False
+    # Rare: Extra arguments
+    get_boolean_choice: Callable[[str], bool] = get_boolean_choice
 
 
 @dataclass
@@ -65,6 +69,7 @@ class LgndrInstallGameArgs:
     disable_https: bool = False
     yes: bool = True
     # Rare: Extra arguments
+    get_boolean_choice: Callable[[str], bool] = get_boolean_choice
     sdl_prompt: Callable[[str, str], List[str]] = lambda a0, a1: []
     verify_stdout: Callable[[int, int, float, float], None] = lambda a0, a1, a2, a3: print(
         f"Verification progress: {a0}/{a1} ({a2:.01f}%) [{a3:.1f} MiB/s]\t\r"

rare/lgndr/api_monkeys.py

@@ -0,0 +1,21 @@
+import logging
+
+from PyQt5.QtWidgets import QMessageBox, QLabel
+
+
+def get_boolean_choice(a0):
+    choice = QMessageBox.question(None, "Import DLCs?", a0)
+    return True if choice == QMessageBox.StandardButton.Yes else False
+
+
+def return_exit(__status):
+    return __status
+
+
+class UILogHandler(logging.Handler):
+    def __init__(self, dest: QLabel):
+        super(UILogHandler, self).__init__()
+        self.widget = dest
+
+    def emit(self, record: logging.LogRecord) -> None:
+        self.widget.setText(record.getMessage())

rare/lgndr/cli.py

@@ -1,7 +1,8 @@
+import functools
 import os
 import logging
 import time
-from typing import Optional, Union
+from typing import Optional, Union, overload
 
 import legendary.cli
 from PyQt5.QtWidgets import QLabel, QMessageBox
@@ -15,25 +16,12 @@ from .core import LegendaryCore
 from .manager import DLManager
 from .api_arguments import LgndrInstallGameArgs, LgndrImportGameArgs, LgndrVerifyGameArgs
 from .api_exception import LgndrException, LgndrLogHandler
-
-
-def get_boolean_choice(message):
-    choice = QMessageBox.question(None, "Import DLCs?", message)
-    return True if choice == QMessageBox.StandardButton.Yes else False
-
-
-class UILogHandler(logging.Handler):
-    def __init__(self, dest: QLabel):
-        super(UILogHandler, self).__init__()
-        self.widget = dest
-
-    def emit(self, record: logging.LogRecord) -> None:
-        self.widget.setText(record.getMessage())
+from .api_monkeys import return_exit, get_boolean_choice
 
 
 class LegendaryCLI(legendary.cli.LegendaryCLI):
     def __init__(self, override_config=None, api_timeout=None):
-        self.core = LegendaryCore(override_config)
+        self.core = LegendaryCore(override_config, timeout=api_timeout)
         self.logger = logging.getLogger('cli')
         self.logging_queue = None
         self.handler = LgndrLogHandler()
@@ -42,16 +30,30 @@ class LegendaryCLI(legendary.cli.LegendaryCLI):
     def resolve_aliases(self, name):
         return super(LegendaryCLI, self)._resolve_aliases(name)
 
-    def prepare_install(self, args: LgndrInstallGameArgs) -> (DLManager, AnalysisResult, InstalledGame, Game, bool, Optional[str], ConditionCheckResult):
-        old_choice = legendary.cli.get_boolean_choice
-        legendary.cli.get_boolean_choice = get_boolean_choice
-        try:
-            return self.install_game(args)
-        except LgndrException as ret:
-            raise ret
-        finally:
-            legendary.cli.get_boolean_choice = old_choice
+    @staticmethod
+    def wrapped(func):
 
+        @functools.wraps(func)
+        def inner(self, args, *oargs, **kwargs):
+            old_exit = legendary.cli.exit
+            legendary.cli.exit = return_exit
+
+            old_choice = legendary.cli.get_boolean_choice
+            if hasattr(args, 'get_boolean_choice') and args.get_boolean_choice is not None:
+                legendary.cli.get_boolean_choice = args.get_boolean_choice
+
+            try:
+                return func(self, args, *oargs, **kwargs)
+            except LgndrException as ret:
+                print(f'Caught exception in wrapped function {ret.message}')
+                raise ret
+            finally:
+                legendary.cli.get_boolean_choice = old_choice
+                legendary.cli.exit = old_exit
+
+        return inner
+
+    @wrapped
     def install_game(self, args: LgndrInstallGameArgs) -> (DLManager, AnalysisResult, InstalledGame, Game, bool, Optional[str], ConditionCheckResult):
         args.app_name = self._resolve_aliases(args.app_name)
         if self.core.is_installed(args.app_name):
@@ -224,6 +226,7 @@ class LegendaryCLI(legendary.cli.LegendaryCLI):
     def _handle_postinstall(self, postinstall, igame, yes=False):
         super(LegendaryCLI, self)._handle_postinstall(postinstall, igame, yes)
 
+    @wrapped
     def uninstall_game(self, args):
         super(LegendaryCLI, self).uninstall_game(args)
 
@@ -336,12 +339,6 @@ class LegendaryCLI(legendary.cli.LegendaryCLI):
                 logger.info(f'Run "legendary repair {args.app_name}" to repair your game installation.')
             return False, len(failed), len(missing)
 
+    @wrapped
     def import_game(self, args: LgndrImportGameArgs):
-        old_choice = legendary.cli.get_boolean_choice
-        legendary.cli.get_boolean_choice = get_boolean_choice
-        try:
-            super(LegendaryCLI, self).import_game(args)
-        except LgndrException as ret:
-            raise ret
-        finally:
-            legendary.cli.get_boolean_choice = old_choice
+        super(LegendaryCLI, self).import_game(args)

rare/shared/__init__.py

@@ -26,7 +26,7 @@ def LegendaryCLISingleton(init: bool = False) -> LegendaryCLI:
     if _legendary_cli_signleton is None and not init:
         raise RuntimeError("Uninitialized use of LegendaryCLISingleton")
     if _legendary_cli_signleton is None:
-        _legendary_cli_signleton = LegendaryCLI()
+        _legendary_cli_signleton = LegendaryCLI(override_config=None, api_timeout=10)
     return _legendary_cli_signleton