Add async recipe fetching to vis

factorygame/data/fetch.py

@@ -6,14 +6,14 @@ from sqlalchemy.orm import Session
 
 from factorygame.data.common import resource_needs_update, chose_resource
 from .models import Base, ResourceFlow, Recipe
-from .sfp import SatisfactoryPlus
+from .sfp import SatisfactoryPlus, DEFAULT_IGNORE_FACTORIES
 from ..helper import click_prompt
 
 
 @click.command()
 @click.option("--debug", is_flag=True)
 @click.option("--refetch", is_flag=True)
-@click.option("--ignore-factories", type=list[str], default=["A.I. Fluid Packer", "Craft Bench", "Equipment Workshop"])
+@click.option("--ignore-factories", type=list[str], default=DEFAULT_IGNORE_FACTORIES)
 @click.argument("search")
 def main(debug: bool, refetch: bool, ignore_factories: list[str], search: str):
     engine = create_engine("sqlite:///file.db", echo=debug)
@@ -28,7 +28,7 @@ def main(debug: bool, refetch: bool, ignore_factories: list[str], search: str):
 
         with SatisfactoryPlus(ignore_factories=ignore_factories, debug=debug) as data_provider:
             if resource is None:
-                ret = data_provider.search_for_resource(session=session, search=search)
+                ret = data_provider.search_for_resource(session=session, search=search, prompt=click_prompt)
                 if ret is None:
                     return
                 else:
@@ -71,6 +71,7 @@ def main(debug: bool, refetch: bool, ignore_factories: list[str], search: str):
             stmt = select(Recipe).distinct().join(Recipe.results).filter(ResourceFlow.resource_id == resource.id)
             for recipe in session.scalars(stmt):
                 print("Result of recipe", recipe.describe())
+        session.commit()
 
 
 if __name__ == "__main__":

factorygame/data/provider.py

@@ -1,4 +1,5 @@
 import abc
+from typing import Callable
 
 from sqlalchemy.orm import Session as AlchemySession
 
@@ -7,7 +8,7 @@ from .models import Resource
 
 class RecipeProvider(abc.ABC):
     @abc.abstractmethod
-    def search_for_resource(self, session: AlchemySession, search: str) -> tuple[Resource, bool]:
+    def search_for_resource(self, session: AlchemySession, search: str, prompt: Callable) -> tuple[Resource, bool]:
         pass
 
     @abc.abstractmethod

factorygame/data/sfp.py

@@ -10,7 +10,8 @@ from sqlalchemy.orm import Session as AlchemySession
 
 from .models import Resource, ResourceFlow, Factory, Recipe
 from .provider import RecipeProvider
-from ..helper import click_prompt
+
+DEFAULT_IGNORE_FACTORIES = ["A.I. Fluid Packer", "Craft Bench", "Equipment Workshop"]
 
 
 class SatisfactoryPlus(RecipeProvider, AbstractContextManager):
@@ -48,7 +49,9 @@ class SatisfactoryPlus(RecipeProvider, AbstractContextManager):
     def __exit__(self, __exc_type, __exc_value, __traceback):
         self._browser_cleanup()
 
-    def search_for_resource(self, session: AlchemySession, search: str) -> tuple[Resource, bool] | None:
+    def search_for_resource(
+        self, session: AlchemySession, search: str, prompt: Callable
+    ) -> tuple[Resource, bool] | None:
         browser = self._init_browser()
         browser.get("https://wiki.kyrium.space/")
         search_bar = browser.find_element(By.CSS_SELECTOR, "nav input[placeholder='Search for an item...']")
@@ -70,7 +73,7 @@ class SatisfactoryPlus(RecipeProvider, AbstractContextManager):
                 options[idx + 1] = name
                 if name.casefold() == search.casefold():
                     default_choice = idx + 1
-            user_choice = click_prompt(options=options, text="Chose a recipe to continue…", default=default_choice)
+            user_choice = prompt(options=options, text="Chose a recipe to continue…", default=default_choice)
             if user_choice is None:
                 return None
             link_html_elem = choices[user_choice - 1]
@@ -93,12 +96,14 @@ class SatisfactoryPlus(RecipeProvider, AbstractContextManager):
         browser.find_element(By.CSS_SELECTOR, "button[id$='tab-0']").click()
         recipes_html_elems = browser.find_elements(By.CSS_SELECTOR, "div[id$='tabpanel-0'] > div > div")
         for recipe_html_elem in recipes_html_elems:
-            self.extract_recipe(session, recipe_html_elem)
+            with session.begin_nested():
+                self.extract_recipe(session, recipe_html_elem)
 
         browser.find_element(By.CSS_SELECTOR, "button[id$='tab-1']").click()
         ingredient_recipes_html_elems = browser.find_elements(By.CSS_SELECTOR, "div[id$='tabpanel-1'] > div > div")
         for recipe_html_elem in ingredient_recipes_html_elems:
-            self.extract_recipe(session, recipe_html_elem)
+            with session.begin_nested():
+                self.extract_recipe(session, recipe_html_elem)
 
         # manual refresh by label, because id might not be set
         updated_resource = session.scalars(Resource.by_label(resource.label)).one()
@@ -114,7 +119,7 @@ class SatisfactoryPlus(RecipeProvider, AbstractContextManager):
             assert factory_label, "factory label is missing (a[text])"
 
             if factory_label.casefold() in self.ignore_factories:
-                return
+                continue
 
             # re-use existing Factory or create new
             factory = session.scalars(Factory.by_label(factory_label)).one_or_none()
@@ -128,13 +133,19 @@ class SatisfactoryPlus(RecipeProvider, AbstractContextManager):
             # ignore recipe when no applicable factories exist
             return
 
+        cached_resources: dict[str, Resource] = {}
+
         def find_or_create_resource(resource_label: str, resource_uri_getter: Callable) -> Resource:
+            if resource_label in cached_resources:
+                return cached_resources[resource_label]
             db_resource = session.scalars(Resource.by_label(resource_label)).one_or_none()
             if db_resource is not None:
+                cached_resources[resource_label] = db_resource
                 return db_resource
 
             resource = Resource(label=resource_label, uri=resource_uri_getter())
             session.add(resource)
+            cached_resources[resource_label] = resource
             return resource
 
         def extract_resource_flow(html_elem):

factorygame/data/vis.py

@@ -2,21 +2,22 @@
 
 import re
 from concurrent.futures import Future
-from datetime import timedelta
 
 import click
+import sqlalchemy
 from NodeGraphQt import BaseNode, NodeBaseWidget
 from NodeGraphQt import NodeGraph, Port
 from NodeGraphQt.constants import PortTypeEnum, NodePropWidgetEnum, ViewerEnum
 from NodeGraphQt.widgets.node_widgets import _NodeGroupBox, NodeLineEdit, NodeCheckBox
 from PySide2 import QtGui
-from PySide2.QtCore import Qt, QObject
+from PySide2.QtCore import Qt, QObject, QThread
-from PySide2.QtWidgets import QSlider, QLineEdit, QCheckBox, QGraphicsItem, QInputDialog
+from PySide2.QtWidgets import QSlider, QLineEdit, QCheckBox, QGraphicsItem, QInputDialog, QDialog, QLabel, QSizePolicy
 from Qt import QtWidgets
 from sqlalchemy import create_engine
 from sqlalchemy.orm import Session
 
 from factorygame.data.common import chose_resource, resource_needs_update, chose_recipe
+from factorygame.data.sfp import SatisfactoryPlus, DEFAULT_IGNORE_FACTORIES
 from .models import Recipe, Resource, ResourceFlow
 
 WORLD_INPUT_PORT_NAME = "World Input"
@@ -30,6 +31,19 @@ OUTPUT_COLOR = (204, 44, 36)
 OTHER_COLOR = (0, 83, 135)
 
 
+def generate_fgbg_stylesheet(bg_alpha: float = 1.0) -> str:
+    dark_bg = QtGui.QColor(*ViewerEnum.BACKGROUND_COLOR.value).darker(120)
+    dark_bg.setAlphaF(bg_alpha)
+    bg_color = dark_bg.getRgb()
+    text_color = tuple(map(str, map(lambda i, j: i - j, (255, 255, 255), bg_color)))
+    return (
+        "* {"
+        "  background-color: rgba(" + ",".join(map(str, bg_color)) + ");"
+        "  color: rgb(" + ",".join(text_color) + ");"
+        "}"
+    )
+
+
 def in_amount_name(label: str) -> str:
     return f"In {label} amount"
 
@@ -77,6 +91,23 @@ def resource_amount_to_text(amount: float):
     return f"{amount:.2f} / min"
 
 
+class AsyncResourceFetcher(QThread):
+    def __init__(self, parent, resource_label: str, engine: sqlalchemy.Engine, debug: bool):
+        super().__init__(parent=parent)
+        self.resource_label = resource_label
+        self.engine = engine
+        self.debug = debug
+
+    def run(self):
+        with Session(self.engine) as session, SatisfactoryPlus(
+            ignore_factories=DEFAULT_IGNORE_FACTORIES, debug=self.debug
+        ) as data_provider:
+            resource = session.scalars(Resource.by_label(self.resource_label)).one()
+            resource = data_provider.update_resource_recipes(session=session, resource=resource)
+            if self.debug:
+                print("Updated resource in separate thread", resource)
+
+
 class NodeSlider(NodeBaseWidget):
     MIN_VALUE = 1
     MAX_VALUE = 250
@@ -378,15 +409,6 @@ class GraphController(QObject):
         super().__init__(parent=parent)
         self.debug = debug
 
-        bg_color = QtGui.QColor(*ViewerEnum.BACKGROUND_COLOR.value).darker(120).getRgb()
-        text_color = tuple(map(str, map(lambda i, j: i - j, (255, 255, 255), bg_color)))
-        self.fgbg_color_stylesheet = (
-            "* {"
-            "  background-color: rgba(" + ",".join(map(str, bg_color)) + ") ;"
-            "  color: rgb(" + ",".join(text_color) + ");"
-            "}"
-        )
-
         self.engine = create_engine("sqlite:///file.db", echo=debug)
         self.graph = NodeGraph(parent=self)
         self.graph.widget.resize(1280, 720)
@@ -415,18 +437,22 @@ class GraphController(QObject):
 
                 resource, exists_in_db = ret
                 if not exists_in_db:
-                    print("Resource not yet fetched, run fetch first")
+                    print("Resource unknown")
                     recipe_selected_future.set_result(None)
                     return
-            if resource_needs_update(resource, recipe_info_timeout=timedelta(days=365)):
+            assert resource is not None
-                print("Please fetch resource", resource.label, "first.")
-                recipe_selected_future.set_result(None)
-                return
 
-            with Session(self.engine) as session:
+            def chose_recipe_async():
-                chose_recipe(session=session, resource=resource, prompt=self.dialog_prompt).add_done_callback(
+                with Session(self.engine) as session:
-                    lambda fut: recipe_selected_future.set_result(fut.result())
+                    chose_recipe(session=session, resource=resource, prompt=self.dialog_prompt).add_done_callback(
-                )
+                        lambda fut: recipe_selected_future.set_result(fut.result())
+                    )
+
+            if resource_needs_update(resource):
+                fetch_future = self.fetch_recipes_async(resource_label=resource.label)
+                fetch_future.add_done_callback(lambda fut: chose_recipe_async)
+            else:
+                chose_recipe_async()
 
         def select_recipe_async():
             with Session(self.engine) as session:
@@ -476,13 +502,21 @@ class GraphController(QObject):
 
                 with Session(self.engine) as session:
                     resource = session.scalars(Resource.by_label(resource_label)).one_or_none()
-                    if resource_needs_update(resource, recipe_info_timeout=timedelta(days=365)):
-                        print("Please fetch resource", resource_label, "first.")
-                        return
 
-                    chose_recipe(session=session, resource=resource, prompt=self.dialog_prompt).add_done_callback(
+                def chose_recipe_async():
-                        recipe_selected_cb
+                    with Session(self.engine) as session:
-                    )
+                        chose_recipe(
+                            session=session,
+                            resource=resource,
+                            prompt=self.dialog_prompt,
+                        ).add_done_callback(recipe_selected_cb)
+
+                if resource_needs_update(resource):
+                    fetch_future = self.fetch_recipes_async(resource_label)
+                    fetch_future.add_done_callback(lambda fut: chose_recipe_async())
+                else:
+                    chose_recipe_async()
+
             elif output_port.name() == WORLD_OUTPUT_PORT_NAME:
                 output_port.clear_connections(push_undo=False, emit_signal=False)
                 for idx, port in enumerate(self.global_input.output_ports()):
@@ -491,7 +525,9 @@ class GraphController(QObject):
                     #     port.connect_to(input_port, push_undo=True, emit_signal=False)
                     #     return
                 if isinstance(input_node, Machine):
-                    resource_amount = parse_resource_amount(str(input_node.get_property(in_amount_name(resource_label))))
+                    resource_amount = parse_resource_amount(
+                        str(input_node.get_property(in_amount_name(resource_label)))
+                    )
                     new_output_port = self.global_input.create_global_output(
                         graph=self.graph,
                         resource_label=resource_label,
@@ -518,13 +554,18 @@ class GraphController(QObject):
 
                     with Session(self.engine) as session:
                         resource = session.scalars(Resource.by_label(resource_label)).one_or_none()
-                        if resource_needs_update(resource, recipe_info_timeout=timedelta(days=365)):
-                            print("Please fetch resource", resource_label, "first.")
-                            return
 
+                    def chose_recipe_async():
                         chose_recipe(
                             session=session, resource=resource, prompt=self.dialog_prompt, ingredient=True
                         ).add_done_callback(recipe_selected_cb)
+
+                    if resource_needs_update(resource):
+                        fetch_future = self.fetch_recipes_async(resource_label=resource_label)
+                        fetch_future.add_done_callback(lambda fut: chose_recipe_async())
+                    else:
+                        chose_recipe_async()
+
                 elif input_port.name() == WORLD_INPUT_PORT_NAME:
                     input_port.clear_connections(push_undo=False, emit_signal=False)
                     for idx, port in enumerate(self.global_output.output_ports()):
@@ -546,12 +587,27 @@ class GraphController(QObject):
                     resource_amount,
                 )
 
+    def fetch_recipes_async(self, resource_label):
+        if self.debug:
+            print("Fetching recipes for resource", resource_label)
+        fetcher_thread = AsyncResourceFetcher(
+            parent=self.graph,
+            resource_label=resource_label,
+            engine=self.engine,
+            debug=self.debug,
+        )
+        fetch_future = Future()
+        self.dialog_loading(fetch_future)
+        fetcher_thread.finished.connect(lambda: fetch_future.set_result(None))
+        fetcher_thread.start()
+        return fetch_future
+
     def dialog_prompt(self, options: dict[int, str], text: str, default: int) -> Future[int | None]:
         if self.debug:
             print("Displaying QInputDialog with options:", ", ".join(options.values()))
 
         dialog = QInputDialog(parent=self.graph.widget)
-        dialog.setStyleSheet(self.fgbg_color_stylesheet)
+        dialog.setStyleSheet(generate_fgbg_stylesheet())
         dialog.setModal(True)
         dialog.setLabelText(text)
         if default in options:
@@ -564,6 +620,21 @@ class GraphController(QObject):
         dialog.show()
         return ret
 
+    def dialog_loading(self, future: Future):
+        dialog = QDialog(parent=self.graph.widget, f=(Qt.Dialog | Qt.ToolTip))
+        dialog.setStyleSheet(generate_fgbg_stylesheet(0.5))
+        dialog.setModal(True)
+        dialog.setSizePolicy(QSizePolicy.Expanding, QSizePolicy.Expanding)
+        label = QLabel("Loading…", dialog)
+        label.setAlignment(Qt.AlignCenter)
+        dialog.show()
+
+        def remove_dialog():
+            dialog.hide()
+            dialog.destroy()
+
+        future.add_done_callback(lambda fut: remove_dialog())
+
 
 @click.command
 @click.option("--debug", is_flag=True)