Refactored duplicated evaluate_task logic for the uis

edomora97 · edomora97 · commit 0d33a24293e2 · 2018-11-17T13:51:22.000+01:00
diff --git a/python/formats/ioi_format.py b/python/formats/ioi_format.py
@@ -267,32 +267,22 @@ def evaluate_task(frontend: Frontend, task: Task, solutions: List[Solution],
         task,
         dict((st_num, [tc for tc in st.testcases.keys()])
              for st_num, st in task.subtasks.items()), config.ui == UIS.PRINT)
+    curses_ui = None
+    finish_ui = None
     if config.ui == UIS.CURSES:
         curses_ui = IOICursesUI(ui_interface)
-        curses_ui.start()
+    if config.ui != UIS.SILENT:
+        finish_ui = IOIFinishUI(config, task, ui_interface)
 
-    try:
+    with ui_interface.run_in_ui(curses_ui, finish_ui):
         ins, outs, vals = generate_inputs(frontend, task, ui_interface, config)
         evaluate_solutions(frontend, ins, outs, vals, solutions, ui_interface,
                            config)
 
         sanity_pre_checks(task, solutions, frontend, config, ui_interface)
         frontend.evaluate()
         sanity_post_checks(task, solutions, ui_interface)
-    except:
-        if config.ui == UIS.CURSES:
-            curses_ui.stop()
-        traceback.print_exc()
-        return ui_interface
-
-    if config.ui == UIS.CURSES:
-        if curses_ui.errored:
-            return ui_interface
-        curses_ui.stop()
 
-    if config.ui != UIS.SILENT:
-        finish_ui = IOIFinishUI(config, task, ui_interface)
-        finish_ui.print()
     return ui_interface
 
 
diff --git a/python/formats/terry_format.py b/python/formats/terry_format.py
@@ -89,11 +89,14 @@ def evaluate_task(frontend: Frontend, task: TerryTask,
                   solutions: List[SourceFile],
                   config: Config) -> TerryUIInterface:
     ui_interface = TerryUIInterface(task, config.ui == UIS.PRINT)
+    curses_ui = None
+    finish_ui = None
     if config.ui == UIS.CURSES:
         curses_ui = TerryCursesUI(ui_interface)
-        curses_ui.start()
+    if config.ui != UIS.SILENT:
+        finish_ui = TerryFinishUI(config, task, ui_interface)
 
-    try:
+    with ui_interface.run_in_ui(curses_ui, finish_ui):
         task.generator.prepare(frontend, config)
         ui_interface.add_non_solution(task.generator)
         if task.validator:
@@ -111,19 +114,6 @@ def evaluate_task(frontend: Frontend, task: TerryTask,
             evaluate_solution(frontend, task, solution, config, ui_interface)
 
         frontend.evaluate()
-    except:
-        if config.ui == UIS.CURSES:
-            curses_ui.stop()
-        traceback.print_exc()
-        return ui_interface
-
-    if config.ui == UIS.CURSES:
-        if curses_ui.errored:
-            return ui_interface
-        curses_ui.stop()
-    if config.ui != UIS.SILENT:
-        finish_ui = TerryFinishUI(config, task, ui_interface)
-        finish_ui.print()
     return ui_interface
 
 
diff --git a/python/uis/__init__.py b/python/uis/__init__.py
@@ -1,11 +1,12 @@
 #!/usr/bin/env python3
 import curses
 from abc import ABC, abstractmethod
+from contextlib import contextmanager
 from enum import Enum
 import signal
 import threading
 import time
-from typing import Dict, List
+from typing import Dict, List, Optional
 import traceback
 
 from task_maker.config import Config
@@ -136,15 +137,37 @@ def add_error(self, message: str):
         self.printer.red("ERROR  ", bold=True)
         self.printer.text(message.strip() + "\n")
 
+    @contextmanager
+    def run_in_ui(self, curses_ui: Optional["CursesUI"],
+                  finish_ui: Optional["FinishUI"]):
+        if curses_ui:
+            curses_ui.start()
+        try:
+            yield
+        except:
+            if curses_ui:
+                curses_ui.stop()
+            traceback.print_exc()
+            return
+        else:
+            if curses_ui:
+                curses_ui.stop()
+        if finish_ui:
+            finish_ui.print()
 
-class FinishUI:
+
+class FinishUI(ABC):
     LIMITS_MARGIN = 0.8
 
     def __init__(self, config: Config, interface: UIInterface):
         self.config = config
         self.interface = interface
         self.printer = StdoutPrinter()
 
+    @abstractmethod
+    def print(self):
+        pass
+
     def _print_compilation(self, solution: str,
                            result: SourceFileCompilationResult,
                            max_sol_len: int):
