SatisfactoryPlusCalculator/factorygame/calculator/SatisfactoryCalculator.py
Ben 9b1cae051a
Add first revision of database interaction
for recipe and information storage.
2024-01-28 19:05:42 +01:00

255 lines
9.7 KiB
Python

#!/usr/bin/env python3
import datetime
from typing import Optional
from urllib.parse import urljoin
import click
import sqlalchemy
from selenium.webdriver import Firefox
from selenium.webdriver.common.by import By
from selenium.webdriver.firefox.options import Options
from selenium.webdriver.remote.webdriver import WebDriver
from sqlalchemy import String, create_engine, select, ForeignKey, Select, Table, Column
from sqlalchemy.orm import DeclarativeBase, Mapped, mapped_column, relationship, Session
class Base(DeclarativeBase):
pass
class Resource(Base):
__tablename__ = "resources"
id: Mapped[int] = mapped_column(primary_key=True)
label: Mapped[str] = mapped_column(String(127))
wiki_url: Mapped[str]
recipes_populated_at: Mapped[Optional[datetime.datetime]]
flows: Mapped[list["ResourceFlow"]] = relationship(back_populates="resource")
@classmethod
def by_label(cls, search: str) -> Select[tuple["Resource"]]:
return select(cls).where(cls.label.ilike(search))
class Factory(Base):
__tablename__ = "factories"
id: Mapped[int] = mapped_column(primary_key=True)
label: Mapped[str] = mapped_column(String(127))
wiki_url: Mapped[str]
@classmethod
def by_label(cls, search: str) -> Select[tuple["Factory"]]:
return select(cls).where(cls.label.ilike(search))
ingredients_table = Table(
"recipe_ingredients",
Base.metadata,
Column("recipe_id", ForeignKey("recipes.id"), primary_key=True),
Column("resource_flow_id", ForeignKey("resource_flows.id"), primary_key=True),
)
results_table = Table(
"recipe_results",
Base.metadata,
Column("recipe_id", ForeignKey("recipes.id"), primary_key=True),
Column("resource_flow_id", ForeignKey("resource_flows.id"), primary_key=True),
)
class ResourceFlow(Base):
__tablename__ = "resource_flows"
id: Mapped[int] = mapped_column(primary_key=True)
ingredient_in: Mapped[Optional["Recipe"]] = relationship(secondary=ingredients_table, back_populates="ingredients")
result_of: Mapped[Optional["Recipe"]] = relationship(secondary=results_table, back_populates="results")
resource_id: Mapped[int] = mapped_column(ForeignKey("resources.id"))
resource: Mapped["Resource"] = relationship(back_populates="flows")
amount: Mapped[str]
time: Mapped[str]
class Recipe(Base):
__tablename__ = "recipes"
id: Mapped[int] = mapped_column(primary_key=True)
factory_id: Mapped[int] = mapped_column(ForeignKey("factories.id"))
factory: Mapped["Factory"] = relationship()
ingredients: Mapped[list["ResourceFlow"]] = relationship(
secondary=ingredients_table, back_populates="ingredient_in"
)
results: Mapped[list["ResourceFlow"]] = relationship(secondary=results_table, back_populates="result_of")
def normalize_url(browser: WebDriver, href: str) -> str:
return urljoin(base=browser.current_url, url=href)
def populate_recipes(browser: WebDriver, engine: sqlalchemy.Engine, input_resource_label: str):
browser.find_element(By.CSS_SELECTOR, "button#\\3Ar1\\3A-tab-0").click()
recipes_html_elems = browser.find_elements(By.CSS_SELECTOR, "#\\3Ar1\\3A-tabpanel-0 > div > div")
for recipe_idx in range(len(recipes_html_elems)):
recipe_html_elem = recipes_html_elems[recipe_idx]
factory_html_elem = recipe_html_elem.find_element(By.CSS_SELECTOR, ".flex-col > span > a")
factory_label = factory_html_elem.text
factory_url = urljoin(base=browser.current_url, url=factory_html_elem.get_attribute("href"))
print("recipe", recipe_idx, "produced in:", factory_label, factory_url)
def extract_resource_flow(html_elem):
resource_img = html_elem.find_element(By.TAG_NAME, "img")
resource_label = resource_img.get_attribute("alt")
wiki_url = normalize_url(
browser=browser,
href=html_elem.find_element(By.TAG_NAME, "a").get_attribute("href"),
)
kwargs = {}
if resource_label == input_resource_label:
kwargs['recipes_populated_at'] = datetime.datetime.utcnow()
resource = Resource(label=resource_label, wiki_url=wiki_url, **kwargs)
amount = html_elem.find_element(By.CSS_SELECTOR, ".text-xs:nth-child(2)").text
time = html_elem.find_element(By.CSS_SELECTOR, ".text-xs:nth-child(3)").text
return ResourceFlow(resource=resource, amount=amount, time=time)
ingredient_html_elems = recipe_html_elem.find_elements(
By.CSS_SELECTOR, f".flex-row > div:nth-child(1) > div:has(> a)"
)
ingredients: list[ResourceFlow] = []
for ingredient_idx in range(len(ingredient_html_elems)):
resource_flow = extract_resource_flow(ingredient_html_elems[ingredient_idx])
ingredients.append(resource_flow)
print(
"recipe",
recipe_idx,
"ingredient",
ingredient_idx,
"name:",
resource_flow.resource.label,
)
print(
"recipe",
recipe_idx,
"ingredient",
ingredient_idx,
"count:",
resource_flow.amount,
)
print(
"recipe",
recipe_idx,
"ingredient",
ingredient_idx,
"time:",
resource_flow.time,
)
result_html_elems = recipe_html_elem.find_elements(
By.CSS_SELECTOR, f".flex-row > div:nth-child(3) > div:has(> a)"
)
results: list[ResourceFlow] = []
for result_idx in range(len(result_html_elems)):
resource_flow = extract_resource_flow(result_html_elems[result_idx])
results.append(resource_flow)
print(
"recipe",
recipe_idx,
"result",
result_idx,
"name:",
resource_flow.resource.label,
)
print(
"recipe",
recipe_idx,
"result",
result_idx,
"count:",
resource_flow.amount,
)
print(
"recipe",
recipe_idx,
"result",
result_idx,
"time:",
resource_flow.time,
)
with Session(engine, autoflush=False) as session:
factory = session.scalars(Factory.by_label(factory_label)).one_or_none()
if not factory:
factory = Factory(label=factory_label, wiki_url=factory_url)
session.add(factory)
for flow in ingredients + results:
res = session.scalars(Resource.by_label(flow.resource.label)).one_or_none()
if res:
flow.resource = res
else:
session.add(flow.resource)
session.add(flow)
recipe = Recipe(factory=factory, ingredients=ingredients, results=results)
session.add(recipe)
session.commit()
@click.command()
@click.option("--result", is_flag=True)
@click.option("--debug", is_flag=True)
@click.argument("search")
def main(result: bool, debug: bool, search: str):
engine = create_engine("sqlite:///file.db", echo=debug)
Base.metadata.create_all(bind=engine)
if result and search:
with Session(engine) as session:
for obj in session.scalars(Resource.by_label(search)):
print(obj)
firefox_options = Options()
if not debug:
firefox_options.add_argument("--headless")
browser = Firefox(options=firefox_options)
try:
browser.get("https://wiki.kyrium.space/")
browser.set_window_size(1600, 1015)
search_bar = browser.find_element(By.CSS_SELECTOR, "nav input[placeholder='Search for an item...']")
search_bar.click()
search_bar.send_keys(search)
search_button = browser.find_element(By.CSS_SELECTOR, "nav button[type='submit']")
search_button.click()
browser.implicitly_wait(5)
choices = browser.find_elements(By.CSS_SELECTOR, "body > div > .container:nth-child(1) a.items-center")
if not choices:
print("No wiki entries found for this result")
return
elif len(choices) > 1:
default_choice = 1
choice_names: list[str] = []
for choice_idx in range(1, len(choices) + 1):
recipe_choice = choices[choice_idx - 1]
name = recipe_choice.find_element(By.TAG_NAME, "img").get_attribute("alt")
choice_names.append(name)
if name.casefold() == search.casefold():
default_choice = choice_idx
print(f"{choice_idx}: {name}")
user_choice = click.prompt("Chose a recipe to continue…", default=default_choice)
if not user_choice:
user_choice = default_choice
else:
user_choice = int(user_choice)
link_html_elem = choices[user_choice - 1]
else:
link_html_elem = choices[0]
resource_label = link_html_elem.find_element(By.TAG_NAME, "img").get_attribute("alt")
# FIXME: check if resource_label is in database
if debug:
print("resource_label:", resource_label)
link_html_elem.click()
populate_recipes(browser=browser, engine=engine, input_resource_label=resource_label)
finally:
if not debug:
browser.quit()