123 lines
4.0 KiB
Python
123 lines
4.0 KiB
Python
import os
|
|
import pandas as pd
|
|
|
|
from craft import CraftComponent, CraftItem, CraftRecipe
|
|
from loguru import logger
|
|
|
|
class CraftStorage:
|
|
def __init__(self):
|
|
# items
|
|
self.items_df = pd.DataFrame(columns=[
|
|
"name", "img_url", "is_base", "img_shift"
|
|
])
|
|
|
|
##Чтобы на конкат не орал
|
|
self.items_df["is_base"] = self.items_df["is_base"].astype('bool')
|
|
|
|
# recipes
|
|
self.recipes_df = pd.DataFrame(columns=[
|
|
"recipe_id", "output_item_name", "output_count", "craft_type"
|
|
])
|
|
|
|
# components
|
|
self.components_df = pd.DataFrame(columns=[
|
|
"recipe_id", "input_item", "count"
|
|
])
|
|
self.recipe_signatures = set()
|
|
self._recipe_id_seq = 1
|
|
|
|
def try_add_recipe_signature(self, output_item_name: str, components: list[CraftComponent], craft_type: str) -> bool:
|
|
# сортируем, чтобы порядок не влиял
|
|
parts = sorted(
|
|
(c.input_item, c.count) for c in components
|
|
)
|
|
sig = (output_item_name, tuple(parts), craft_type)
|
|
if sig in self.recipe_signatures:
|
|
logger.warning(f'Дубликат рецепта {output_item_name}. Пропускаю ')
|
|
#print(f'CRAFT: {}')
|
|
# for el in components:
|
|
# print(f'{el.input_item} : {el.count}')
|
|
# print('Пропускаю')
|
|
# print('....')
|
|
return False
|
|
|
|
self.recipe_signatures.add(sig)
|
|
return True
|
|
|
|
|
|
def add_item(self, item: CraftItem):
|
|
# проверяем, есть ли уже такой item
|
|
if item.name in self.items_df["name"].values:
|
|
logger.warning(f'{item.name} уже есть в датасете, скипаю')
|
|
return # уже есть, не дублируем
|
|
|
|
row = {
|
|
"name": item.name,
|
|
"img_url": item.img_url,
|
|
"is_base": bool(item.is_base),
|
|
"img_shift": item.img_shift
|
|
}
|
|
|
|
self.items_df = pd.concat(
|
|
[self.items_df, pd.DataFrame([row])],
|
|
ignore_index=True
|
|
)
|
|
|
|
|
|
def add_recipe(self, recipe: CraftRecipe) -> int:
|
|
recipe_id = self._recipe_id_seq
|
|
self._recipe_id_seq += 1
|
|
|
|
recipe.recipe_id = recipe_id
|
|
|
|
row = {
|
|
"recipe_id": recipe_id,
|
|
"output_item_name": recipe.output_item_name,
|
|
"output_count": recipe.output_count,
|
|
"craft_type": recipe.craft_type
|
|
}
|
|
|
|
self.recipes_df = pd.concat(
|
|
[self.recipes_df, pd.DataFrame([row])],
|
|
ignore_index=True
|
|
)
|
|
|
|
return recipe_id
|
|
|
|
|
|
|
|
def add_component(self, component: CraftComponent):
|
|
# защита от дубликатов
|
|
exists = (
|
|
(self.components_df["recipe_id"] == component.recipe_id) &
|
|
(self.components_df["input_item"] == component.input_item) &
|
|
(self.components_df['count'] == component.count)
|
|
).any()
|
|
|
|
if exists:
|
|
logger.warning(f"""
|
|
Уже существует такая запись для рецепта: {component.recipe_id} <--- {component.input_item}*{component.count}
|
|
Скипаем
|
|
""")
|
|
return
|
|
|
|
row = {
|
|
"recipe_id": component.recipe_id,
|
|
"input_item": component.input_item,
|
|
"count": component.count
|
|
}
|
|
|
|
self.components_df = pd.concat(
|
|
[self.components_df, pd.DataFrame([row])],
|
|
ignore_index=True
|
|
)
|
|
|
|
def save(self):
|
|
os.makedirs('data', exist_ok=True)
|
|
self.items_df.to_csv('data/items.csv', index=False)
|
|
self.recipes_df.to_csv('data/recipes.csv', index=False)
|
|
self.components_df.to_csv('data/recipe_components.csv', index=False)
|
|
self.items_df.to_excel('data/items.xlsx', index=False)
|
|
self.recipes_df.to_excel('data/recipes.xlsx', index=False)
|
|
self.components_df.to_excel('data/recipe_components.xlsx', index=False)
|