models.py

2.3 kB · python · 75 lines

1from dataclasses import dataclass, field2from typing import Any, Dict, List3from mrlypy.paint.enums import Edition, Ink4from mrlypy.paint.models import Paint5from mrlypy.tile.models import Tile6from PIL import Image78# FILE910@dataclass11class File:1213    width: int = None14    height: int = None15    data: Image.Image = None1617    def to_dict(self) -> Dict[str, Any]:18        return {19            "width": self.width,20            "height": self.height,21        }2223    @classmethod24    def from_dict(cls, data: Dict[str, Any]) -> "File":25        f = cls()26        f.width = data["width"]27        f.height = data["height"]28        return f2930# GEN3132@dataclass33class Gen:3435    key: str = None36    seed: int = None37    edition: Edition = None38    primaries: List[Ink] = None39    tile: Tile = None40    mask: Tile = None41    paint: Paint = None42    files: List[File] = field(default_factory=list)4344    @property45    def is_cover(self) -> bool:46        return self.edition in [Edition.ROWS, Edition.COLUMNS, Edition.RANDOM]4748    @property49    def is_prime(self) -> bool:50        return self.edition in [Edition.LAYERS, Edition.NEIGHBORS]5152    def to_dict(self) -> Dict[str, Any]:53        data = {}54        data["key"] = self.key55        data["seed"] = self.seed56        data["edition"] = self.edition.value if self.edition else None57        data["primaries"] = [p.value for p in self.primaries] if self.primaries else None58        data["tile"] = self.tile.to_dict() if self.tile else None59        data["mask"] = self.mask.to_dict() if self.mask else None60        data["paint"] = self.paint.to_dict() if self.paint else None61        data["files"] = [f.to_dict() for f in self.files]62        return data6364    @classmethod65    def from_dict(cls, data: Dict[str, Any]) -> "Gen":66        v = cls()67        v.key = data["key"]68        v.seed = data["seed"]69        v.edition = Edition(data["edition"]) if data.get("edition") else None70        v.primaries = [Ink(p) for p in data["primaries"]] if data.get("primaries") else None71        v.tile = Tile.from_dict(data["tile"]) if data.get("tile") else None72        v.mask = Tile.from_dict(data["mask"]) if data.get("mask") else None73        v.paint = Paint.from_dict(data["paint"]) if data.get("paint") else None74        v.files = [File.from_dict(f) for f in data.get("files", [])]75        return v