Add async recipe fetching to vis

This commit is contained in:
Ben 2024-02-02 21:48:57 +01:00
parent 9f833c33fe
commit ce0459c832
Signed by: ben
GPG Key ID: 0F54A7ED232D3319
4 changed files with 126 additions and 42 deletions

View File

@ -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__":

View File

@ -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

View File

@ -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):

View File

@ -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.QtWidgets import QSlider, QLineEdit, QCheckBox, QGraphicsItem, QInputDialog
from PySide2.QtCore import Qt, QObject, QThread
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)):
print("Please fetch resource", resource.label, "first.")
recipe_selected_future.set_result(None)
return
assert resource is not None
with Session(self.engine) as session:
chose_recipe(session=session, resource=resource, prompt=self.dialog_prompt).add_done_callback(
lambda fut: recipe_selected_future.set_result(fut.result())
)
def chose_recipe_async():
with Session(self.engine) as session:
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(
recipe_selected_cb
)
def chose_recipe_async():
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)