Refactor, deduplicate some code

factorygame/data/common.py

@@ -0,0 +1,48 @@
+from datetime import timedelta, datetime
+from typing import Callable, Optional
+
+import sqlalchemy
+from sqlalchemy import select
+
+from .models import Resource, Recipe, ResourceFlow
+
+
+def resource_needs_update(resource: Resource | None, recipe_info_timeout: Optional[timedelta] = None) -> bool:
+    if recipe_info_timeout is None:
+        recipe_info_timeout = timedelta(days=30)
+    return (
+        resource is None
+        or resource.recipes_populated_at is None
+        or datetime.utcnow() - resource.recipes_populated_at > recipe_info_timeout
+    )
+
+
+def chose_resource(session: sqlalchemy.Session, resource_label: str, prompt: Callable) -> Resource | None:
+    matching_resources = session.scalars(Resource.by_label(resource_label)).all()
+    if len(matching_resources) == 0:
+        print("Could not find existing resources matching the search string.. starting wiki search")
+    else:
+        options = {(idx + 1): str(matching_resources[idx].label) for idx in range(len(matching_resources))}
+        options[0] = ""
+        selected_res = prompt(
+            options=options, text="Chose a resource to continue or 0 to continue with a wiki search", default=1
+        )
+        if selected_res is not None and selected_res != 0:
+            return matching_resources[selected_res - 1]
+    return None
+
+
+def chose_recipe(session: sqlalchemy.Session, resource: Resource, prompt: Callable) -> Recipe | None:
+    stmt = select(Recipe).join(Recipe.results).filter(ResourceFlow.resource_id == resource.id)
+    recipes = session.scalars(stmt).all()
+    if len(recipes) == 0:
+        print("No recipes found for resource")
+        return None
+    elif len(recipes) > 1:
+        options = {(idx + 1): recipes[idx].describe() for idx in range(len(recipes))}
+        user_choice = prompt(options=options, text="Select recipe", default=1)
+        if user_choice is None:
+            return None
+        return recipes[user_choice - 1]
+    else:
+        return recipes[0]

factorygame/data/fetch.py

@@ -1,18 +1,14 @@
 #!/usr/bin/env python3
 
-import datetime
-from typing import Optional
-
 import click
 from sqlalchemy import create_engine, select
 from sqlalchemy.orm import Session
 
-from .models import Base, Resource, ResourceFlow, Recipe
+from factorygame.data.common import resource_needs_update, chose_resource
+from .models import Base, ResourceFlow, Recipe
 from .sfp import SatisfactoryPlus
 from ..helper import prompt
 
-__recipe_info_timeout = datetime.timedelta(days=30)
-
 
 @click.command()
 @click.option("--result", is_flag=True)
@@ -23,53 +19,32 @@ def main(result: bool, debug: bool, refetch: bool, search: str):
     engine = create_engine("sqlite:///file.db", echo=debug)
     Base.metadata.create_all(bind=engine)
     if result and search:
-        do_provider_search = True
-        resource: Optional[Resource] = None
         with Session(engine) as session:
-            matching_resources = session.scalars(Resource.by_label(search)).all()
+            resource = chose_resource(session=session, resource_label=search, prompt=prompt)
-            if len(matching_resources) == 0:
+            exists_in_db = resource is not None
-                print("Could not find existing resources matching the search string.. starting wiki search")
-            else:
-                options = {(idx + 1): str(matching_resources[idx].label) for idx in range(len(matching_resources))}
-                options[0] = ""
-                selected_res = prompt(
-                    options=options, text="Chose a resource to continue or 0 to continue with a wiki search", default=1
-                )
-                if selected_res is None:
-                    return
-                if selected_res != 0:
-                    resource = matching_resources[selected_res - 1]
-                    do_provider_search = False
-                    exists_in_db = True
 
             with SatisfactoryPlus(debug=debug) as data_provider:
-                if do_provider_search:
+                if resource is None:
                     ret = data_provider.search_for_resource(session=session, search=search)
                     if ret is None:
                         return
                     else:
                         resource, exists_in_db = ret
 
-                refetch = (
+                refetch |= resource_needs_update(resource)
-                    refetch
-                    or resource.recipes_populated_at is None
-                    or datetime.datetime.utcnow() - resource.recipes_populated_at > __recipe_info_timeout
-                )
-                if refetch and exists_in_db:
-                    print("Deleting recipes for", resource.label)
-                    with session.begin_nested():
-                        for flow in session.scalars(
-                            select(ResourceFlow).where(ResourceFlow.resource_id == resource.id)
-                        ):
-                            if flow.result_of:
-                                for flow2 in flow.result_of.ingredients:
-                                    session.delete(flow2)
-                                for flow2 in flow.result_of.results:
-                                    session.delete(flow2)
-                                session.delete(flow.result_of)
-
                 if refetch:
                     if exists_in_db:
+                        print("Deleting recipes for", resource.label)
+                        with session.begin_nested():
+                            for flow in session.scalars(
+                                select(ResourceFlow).where(ResourceFlow.resource_id == resource.id)
+                            ):
+                                if flow.result_of:
+                                    for flow2 in flow.result_of.ingredients:
+                                        session.delete(flow2)
+                                    for flow2 in flow.result_of.results:
+                                        session.delete(flow2)
+                                    session.delete(flow.result_of)
                         print("Refetching recipes for", resource.label)
                     else:
                         print("Fetching recipes for new resource", resource.label)

factorygame/data/provider.py

@@ -1,6 +1,6 @@
 import abc
 
-from sqlalchemy.orm import Session
+import sqlalchemy
 
 from .models import Resource
 
@@ -8,9 +8,9 @@ from .models import Resource
 class RecipeProvider(abc.ABC):
 
     @abc.abstractmethod
-    def search_for_resource(self, session: Session, search: str) -> tuple[Resource, bool]:
+    def search_for_resource(self, session: sqlalchemy.Session, search: str) -> tuple[Resource, bool]:
         pass
 
     @abc.abstractmethod
-    def update_resource_recipes(self, session: Session, resource: Resource):
+    def update_resource_recipes(self, session: sqlalchemy.Session, resource: Resource):
         pass

factorygame/data/sfp.py

@@ -3,10 +3,10 @@ from datetime import datetime
 from typing import Optional
 from urllib.parse import urljoin
 
+import sqlalchemy
 from selenium.webdriver import Firefox
 from selenium.webdriver.common.by import By
 from selenium.webdriver.firefox.options import Options
-from sqlalchemy.orm import Session
 
 from .models import Resource, ResourceFlow, Factory, Recipe
 from .provider import RecipeProvider
@@ -44,7 +44,7 @@ class SatisfactoryPlus(RecipeProvider, AbstractContextManager):
     def __exit__(self, __exc_type, __exc_value, __traceback):
         self._browser_cleanup()
 
-    def search_for_resource(self, session: Session, search: str) -> tuple[Resource, bool] | None:
+    def search_for_resource(self, session: sqlalchemy.Session, search: str) -> 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...']")
@@ -81,7 +81,7 @@ class SatisfactoryPlus(RecipeProvider, AbstractContextManager):
             resource_fetch_url = self._normalize_url(href=link_html_elem.get_attribute("href"))
             return Resource(label=alt_resource_label, wiki_url=resource_fetch_url), False
 
-    def update_resource_recipes(self, session: Session, resource: Resource) -> Resource:
+    def update_resource_recipes(self, session: sqlalchemy.Session, resource: Resource) -> Resource:
         assert resource.wiki_url, "Resource wiki url not set"
         browser = self._init_browser()
         browser.get(resource.wiki_url)

factorygame/data/vis.py

@@ -1,5 +1,7 @@
 #!/usr/bin/env python3
 
+from datetime import timedelta
+
 import click
 from NodeGraphQt import BaseNode, NodeBaseWidget
 from NodeGraphQt import NodeGraph, Port
@@ -7,11 +9,11 @@ from NodeGraphQt.constants import PortTypeEnum
 from PySide2.QtCore import Qt
 from PySide2.QtWidgets import QSlider
 from Qt import QtWidgets
-from sqlalchemy import create_engine, select
+from sqlalchemy import create_engine
 from sqlalchemy.orm import Session
 
-from .models import Recipe, ResourceFlow, Resource
+from factorygame.data.common import chose_resource, resource_needs_update, chose_recipe
-from .sfp import SatisfactoryPlus
+from .models import Recipe, Resource
 from ..helper import prompt
 
 
@@ -127,50 +129,17 @@ def on_port_connected(input_port: Port, output_port: Port):
         print(f"Port {output_port} connected to {input_port}")
     if isinstance(output_port.node(), GlobalInput) and output_port.name() == "Create Machine":
         output_port.clear_connections(push_undo=False, emit_signal=False)
-        with Session(engine) as session, SatisfactoryPlus() as data_provider:
+        with Session(engine) as session:
-            do_provider_search = False
             resource_label = input_port.name()
-            matching_resources = session.scalars(Resource.by_label(resource_label)).all()
+            resource = session.scalars(Resource.by_label(resource_label)).one_or_none()
-            if len(matching_resources) == 0:
+            if resource_needs_update(resource, recipe_info_timeout=timedelta(days=365)):
-                print("Could not find existing resources matching the search string.. starting wiki search")
+                print("Please fetch resource", resource_label, "first.")
-            else:
+                return
-                options = {(idx + 1): str(matching_resources[idx].label) for idx in range(len(matching_resources))}
+
-                options[0] = ""
+            # FIXME: use some Qt UI prompt method
-                selected_res = prompt(
+            recipe = chose_recipe(session=session, resource=resource, prompt=prompt)
-                    options=options,
+            if recipe is None:
-                    text="Chose a resource to continue or 0 to continue with a wiki search",
-                    default=1,
-                )
-                if selected_res is None:
-                    return
-                if selected_res != 0:
-                    resource = matching_resources[selected_res - 1]
-                    do_provider_search = False
-
-            if do_provider_search:
-                ret = data_provider.search_for_resource(session=session, search=resource_label)
-                if ret is None:
-                    print("Resource not found")
-                    return
-
-                resource, exists_in_db = ret
-                if not exists_in_db:
-                    print("Resource not yet fetched, run fetch first")
-                    return
-
-            stmt = select(Recipe).join(Recipe.results).filter(ResourceFlow.resource_id == resource.id)
-            recipes = session.scalars(stmt).all()
-            if len(recipes) == 0:
-                print("No recipes found for resource")
                 return
-            elif len(recipes) > 1:
-                options = {(idx + 1): recipes[idx].describe() for idx in range(len(recipes))}
-                user_choice = prompt(options=options, text="Select recipe", default=1)
-                if user_choice is None:
-                    return
-                recipe = recipes[user_choice - 1]
-            else:
-                recipe = recipes[0]
 
         recipe_machine = graph.create_node("factorygame.Machine", push_undo=True)
         recipe_machine.assign_recipe(recipe)
@@ -188,25 +157,14 @@ def main(debug: bool, search: str):
     global engine
     globals()["debug"] = debug
     engine = create_engine("sqlite:///file.db", echo=debug)
-    with Session(engine) as session, SatisfactoryPlus(debug=debug) as data_provider:
+    with Session(engine) as session:
-        do_provider_search = False
+        # FIXME: use some Qt UI prompt method
-        matching_resources = session.scalars(Resource.by_label(search)).all()
+        resource = chose_resource(session=session, resource_label=search, prompt=prompt)
-        if len(matching_resources) == 0:
-            print("Could not find existing resources matching the search string.. starting wiki search")
-        else:
-            options = {(idx + 1): str(matching_resources[idx].label) for idx in range(len(matching_resources))}
-            options[0] = ""
-            selected_res = prompt(
-                options=options, text="Chose a resource to continue or 0 to continue with a wiki search", default=1
-            )
-            if selected_res is None:
-                return
-            if selected_res != 0:
-                resource = matching_resources[selected_res - 1]
-                do_provider_search = False
 
-        if do_provider_search:
+        if resource is None:
-            ret = data_provider.search_for_resource(session=session, search=search)
+            # FIXME: use concurrent resource searching/fetching
+            # ret = data_provider.search_for_resource(session=session, search=search)
+            ret = None
             if ret is None:
                 print("Resource not found")
                 return
@@ -215,20 +173,14 @@ def main(debug: bool, search: str):
             if not exists_in_db:
                 print("Resource not yet fetched, run fetch first")
                 return
-
+        if resource_needs_update(resource, recipe_info_timeout=timedelta(days=365)):
-        stmt = select(Recipe).join(Recipe.results).filter(ResourceFlow.resource_id == resource.id)
+            print("Please fetch resource", resource.label, "first.")
-        recipes = session.scalars(stmt).all()
+            return
-        if len(recipes) == 0:
+
-            print("No recipes found for resource")
+        # FIXME: use some Qt UI prompt method
+        recipe = chose_recipe(session=session, resource=resource, prompt=prompt)
+        if recipe is None:
             return
-        elif len(recipes) > 1:
-            options = {(idx + 1): recipes[idx].describe() for idx in range(len(recipes))}
-            user_choice = prompt(options=options, text="Select recipe", default=1)
-            if user_choice is None:
-                return
-            recipe = recipes[user_choice - 1]
-        else:
-            recipe = recipes[0]
 
     app = QtWidgets.QApplication([])
     global graph