world_data/models.py at main · silasary/world_data · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
import enum
from functools import reduce
import os
import pathlib

import attrs
import re
import json

abspath = os.path.dirname(__file__)
world_folder = pathlib.Path(abspath, "worlds")

class ItemClassification(enum.Flag):
    unknown = 0
    trap = 1
    filler = 2
    useful = 4
    progression = 8
    mcguffin = 16

    bad_name = 256

    @staticmethod
    def from_network_flag(flag: int) -> "ItemClassification":
        if flag == 0:
            return ItemClassification.filler
        value = ItemClassification.unknown
        if flag & 0b00001:
            value |= ItemClassification.progression
        if flag & 0b00010:
            value |= ItemClassification.useful
        if flag & 0b00100:
            value |= ItemClassification.trap
        # if flag & 0b01000:
        #     value |= ItemClassification.mcguffin  # skip balancing
        # if flag & 0b10000:
        #     value |= ItemClassification.mcguffin  # deprioritized
        if value == ItemClassification.progression | ItemClassification.useful:
            value = ItemClassification.progression  # Useful + Progression is just Progression
        return value

classifications = {v.name: v for v in ItemClassification.__members__.values()}

@attrs.define()
class Datapackage:
    items: dict[str, ItemClassification] = attrs.field(factory=dict)
    categories: dict[str, ItemClassification] = attrs.field(factory=dict)
    game_name: str = attrs.field(default="")

    def icon(self, item_name: str) -> str:
        classification = self.items.get(item_name, ItemClassification.unknown)
        emoji = "❓"
        if classification == ItemClassification.mcguffin:
            emoji = "✨"
        if classification == ItemClassification.filler:
            emoji = "<:filler:1277502385459171338>"
        if classification == ItemClassification.useful:
            emoji = "<:useful:1277502389729103913>"
        if classification == ItemClassification.progression:
            emoji = "<:progression:1277502382682542143>"
        if classification == ItemClassification.trap:
            emoji = "❌"
        return emoji

    def set_classification(self, item_name: str, classification: ItemClassification) -> bool:
        if classification == ItemClassification.unknown and self.items.get(item_name, ItemClassification.unknown) != ItemClassification.unknown:
            # We don't want to set an item to unknown if it's already classified
            return False
        if self.items.get(item_name) == classification:
            return False
        self.items[item_name] = classification
        return True

    def postprocess_item_classification(self, item_name: str, classification: ItemClassification) -> ItemClassification:
        """Hook for processing classification that came from the API."""
        # This is mostly for applying mcguffins
        disk_classification = self.items.get(item_name)
        if disk_classification is None or disk_classification == ItemClassification.unknown:
            self.set_classification(item_name, classification)
        if disk_classification == ItemClassification.mcguffin:
            classification = ItemClassification.mcguffin
        return classification


def load_datapackage(game_name, dp: Datapackage = None, follow_redirect: bool = True) -> Datapackage:
    if dp is None:
        dp = Datapackage()
    game_name = game_name.replace("/", "_").replace(":", "_")
    if os.path.exists(world_folder.joinpath(game_name, "redirect.txt")) and follow_redirect:
        game_name = world_folder.joinpath(game_name, "redirect.txt").read_text().strip()
    dp.game_name = game_name

    info_json = world_folder.joinpath(game_name, "progression.json")
    if info_json.exists():
        info = json.loads(info_json.read_text(encoding="utf-8"))
        if 'items' in info:
            for name, classification in info['items'].items():
                v = reduce((lambda a, b: a | b), {ItemClassification[str_classification] for str_classification in classification.split("|")})
                dp.set_classification(name, v)

    prog_txt = world_folder.joinpath(game_name, "progression.txt")
    if prog_txt.exists():
        try:
            progressionFile = open(prog_txt, "r", encoding="utf-8", errors="backslashreplace")
            text = progressionFile.read()
            progressionFile.close()
        except UnicodeDecodeError as e:
            print(f"Error reading {prog_txt}: {e}")
            raise

        for x in text.splitlines():
            if not x:
                continue
            match = re.match(r"^(.*): (.*)$", x)
            classification = reduce((lambda a, b: a | b), {ItemClassification[str_classification] for str_classification in match[2].strip().split("|")})
            dp.set_classification(match[1], classification)

    categories_txt = world_folder.joinpath(game_name, "categories.txt")
    if categories_txt.exists():
        categoriesFile = open(categories_txt, "r", encoding="utf-8")
        text = categoriesFile.read()
        categoriesFile.close()
        for x in text.splitlines():
            if not x:
                continue
            match = re.match(r"^(.*): (.*)$", x)
            classification = reduce((lambda a, b: a | b), {ItemClassification[str_classification] for str_classification in match[2].strip().split("|")})
            dp.categories[match[1]] = classification

    return dp

def save_datapackage(game_name, dp: Datapackage) -> Datapackage:
    game_name = game_name.replace("/", "_").replace(":", "_")
    if os.path.exists(world_folder.joinpath(game_name, "redirect.txt")):
        game_name = world_folder.joinpath(game_name, "redirect.txt").read_text().strip()
    info = {}
    for name, classification in dp.items.items():
        info[name] = classification.name

    world_folder.joinpath(game_name).mkdir(parents=True, exist_ok=True)

    if dp.categories or world_folder.joinpath(game_name, "progression.json").exists():
        save_complex(game_name, dp, info)
        return dp

    progressionFile = world_folder.joinpath(game_name, "progression.txt")
    lines = [f"{k}: {v.name}" for k, v in dp.items.items()]
    lines.sort()
    progressionFile.write_text("\n".join(lines) + "\n", encoding="utf-8")

    return dp

def save_complex(game_name, dp: Datapackage, info: dict[str, str]) -> None:
    dump = json.dumps({"categories": dp.categories, "items": info}, indent=2, sort_keys=True)
    world_folder.joinpath(game_name, "progression.json").write_text(dump, encoding="utf-8")
    progressionFile = world_folder.joinpath(game_name, "progression.txt")
    if progressionFile.exists():
        progressionFile.unlink()