ai_sandbox/data.py

86 lines
2.3 KiB
Python

from pydantic import BaseModel
from recipe_scrapers._abstract import AbstractScraper
from ingredient_parser.postprocess import ParsedIngredient
class Recipe(BaseModel):
id: int | None = None
title: str | None
url: str | None
site_name: str
cuisine: str
category: str
description: str
cook_time: float
prep_time: float
yields: str
class Ingredient(BaseModel):
id: int | None = None
recipe_id: int | None = None
text: str | None = None
purpose: str | None = None
name: str | None
name_conf: float | None
amount_quantity: list[str] | None
amount_unit: list[str] | None
amount_conf: list[float] | None
preparation: str | None
preparation_conf: float | None
comment: str | None
comment_conf: float | None
def parsed_ingredient_to_ingredient(
parsed_ingredient: ParsedIngredient,
purpose=None,
recipe_id=None,
) -> Ingredient:
data = {}
data["purpose"] = purpose
data["recipe_id"] = recipe_id
fields = ["name", "preparation", "comment"]
data["text"] = parsed_ingredient.sentence
for field in fields:
attr = getattr(parsed_ingredient, field)
data[field] = attr.text if attr else None
data[field + "_conf"] = attr.confidence if attr else None
amounts = parsed_ingredient.amount or []
if amounts:
data["amount_quantity"] = []
data["amount_unit"] = []
data["amount_conf"] = []
for amount in amounts:
data["amount_quantity"].append(amount.quantity if amount else "")
data["amount_unit"].append(amount.unit if amount else "")
data["amount_conf"].append(amount.confidence if amount else 0)
else:
data["amount_quantity"] = None
data["amount_unit"] = None
data["amount_conf"] = None
return Ingredient(**data)
def scraped_reciepe_to_recipe(recipe: AbstractScraper) -> Recipe:
site_name = recipe.site_name()
if not site_name:
site_name = ""
site_name = str(site_name)
return Recipe(
title=recipe.title(),
site_name=site_name,
url=recipe.url,
cuisine=recipe.cuisine(),
category=recipe.category(),
description=recipe.description(),
cook_time=recipe.cook_time(),
prep_time=recipe.prep_time(),
yields=recipe.yields(),
)