create.py

5.7 kB · python · 185 lines

1import mrlypy.gen2import mrlypy.paint3import mrlypy.tile4import random5import time6from automator.core.api import logger7from automator.core.errors import NoTaskError8from automator.core.models import Mockup, Placement, Printfile, Task, Variant9from automator.core.s3 import load_paths, load_product, save_paths, task_exists10from automator.core.steps import Step11from mrlypy.core.helpers import hex_key12from mrlypy.paint.colors import get_primary_inks13from typing import Any1415MAX_VARIANTS = 10016MAX_PATHS = 2017PRINTFILE = "printfile"1819def create_task() -> Task:20    task = Task()21    task.created_at = int(time.time())22    while True:23        key = hex_key()24        if not task_exists(key):25            task.key = key26            return task2728def choose_path(task: Task) -> Task:29    paths = load_paths()30    open_products = [product for product, state in paths.items() if state is True]31    if not open_products:32        for product, state in paths.items():33            if state is False:34                paths[product] = True35        open_products = [product for product, state in paths.items() if state is True]36    if not open_products:37        raise NoTaskError("every path is quarantined")38    product = random.choice(open_products)39    paths[product] = False40    save_paths(paths)41    task.product.id = int(product)42    return task4344def quarantine(task: Task) -> None:45    paths = load_paths()46    paths[str(task.product.id)] = None47    save_paths(paths)4849def parse_basics(task: Task, data: dict) -> Task:50    task.product.id = data["id"]51    task.product.category = data["category"]52    task.product.title = data["title"]53    task.product.technique = data["technique"]54    task.product.primaries = data["primaries"]55    task.product.stitch_colors = data["stitch_colors"]56    return task5758def parse_placements(task: Task, data: dict) -> Task:59    task.placements = [60        Placement(61            name=placement["name"],62            width=placement["width"],63            height=placement["height"],64            dpi=placement["dpi"],65        )66        for placement in data["placements"]67        if not placement["is_ignored"]68    ]69    return task7071def create_keys(items: list[Any], task_key: str) -> list[Any]:72    seen: set[str] = set()73    for item in items:74        while True:75            key = hex_key()76            if key not in seen:77                seen.add(key)78                break79        item.key = key80        item.name = f"{task_key}-{key}"81    return items8283def dedup_variants(task: Task) -> Task:84    kept: dict[str, Variant] = {}85    for variant in task.variants:86        first = kept.get(variant.size)87        if first is None:88            kept[variant.size] = variant89            continue90        if first.color != variant.color:91            logger.info(f"{task.desc} dropped variant {variant.id} size {variant.size} colour {variant.color}")92    task.variants = list(kept.values())[:MAX_VARIANTS]93    return task9495def parse_variants(task: Task, data: dict) -> Task:96    task.variants = [97        Variant(98            id=variant["id"],99            cost=variant["cost"],100            size=variant["size"],101            color=variant["color"],102        )103        for variant in data["variants"]104        if not variant["is_ignored"]105    ]106    task = dedup_variants(task)107    task.variants = create_keys(task.variants, task.key)108    return task109110def parse_mockups(task: Task, data: dict) -> Task:111    task.mockups = [112        Mockup(113            id=mockup["id"],114            category=mockup["category"],115            title=mockup["title"],116            variant_ids=mockup["variant_ids"],117        )118        for mockup in data["mockups"]119        if not mockup["is_ignored"]120    ]121    task.mockups = create_keys(task.mockups, task.key)122    return task123124def name_printfiles(printfiles: list[Printfile]) -> list[Printfile]:125    if len(printfiles) == 1:126        printfiles[0].name = PRINTFILE127        return printfiles128    for i, printfile in enumerate(printfiles, 1):129        printfile.name = f"{PRINTFILE}-{i}"130    return printfiles131132def parse_printfiles(task: Task) -> Task:133    printfiles = {}134    for placement in task.placements:135        if placement.id in printfiles:136            continue137        printfiles[placement.id] = Printfile(138            id=placement.id,139            width=placement.width,140            height=placement.height,141            dpi=placement.dpi,142        )143    task.printfiles = list(printfiles.values())144    task.printfiles = create_keys(task.printfiles, task.key)145    task.printfiles = name_printfiles(task.printfiles)146    return task147148def parse_product(task: Task) -> Task:149    data = load_product(task.product.id)150    task = parse_basics(task, data)151    task = parse_placements(task, data)152    task = parse_variants(task, data)153    task = parse_mockups(task, data)154    task = parse_printfiles(task)155    return task156157def create_variation(task: Task) -> Task:158    primaries = get_primary_inks(task.product.primaries)159    config = mrlypy.gen.Config(160        tile=mrlypy.tile.Config(min_size=3, max_size=9, anti=False),161        paint=mrlypy.paint.Config(primaries=primaries),162        files=[],163    )164    gen = mrlypy.gen.create(config)165    gen.key = task.key166    task.seed = gen.seed167    task.variation = gen.to_dict()168    return task169170def open_product(task: Task) -> Task:171    for _ in range(MAX_PATHS):172        task = choose_path(task)173        try:174            return parse_product(task)175        except NoTaskError:176            quarantine(task)177            logger.info(f"quarantined product {task.product.id}, no product json in the bucket")178    raise NoTaskError(f"no product json under {MAX_PATHS} open paths")179180def mrly_create() -> Task:181    task = create_task()182    task = open_product(task)183    task = create_variation(task)184    task.place(Step.GENERATE)185    return task