diff --git a/worlds/lingo/options.py b/worlds/lingo/options.py index 05fb4ed9..65f27269 100644 --- a/worlds/lingo/options.py +++ b/worlds/lingo/options.py @@ -3,7 +3,7 @@ from dataclasses import dataclass from schema import And, Schema from Options import Toggle, Choice, DefaultOnToggle, Range, PerGameCommonOptions, StartInventoryPool, OptionDict -from worlds.lingo.items import TRAP_ITEMS +from .items import TRAP_ITEMS class ShuffleDoors(Choice): diff --git a/worlds/lingo/static_logic.py b/worlds/lingo/static_logic.py index c7ee0010..ff820dd0 100644 --- a/worlds/lingo/static_logic.py +++ b/worlds/lingo/static_logic.py @@ -78,13 +78,16 @@ def get_progressive_item_id(name: str): def load_static_data_from_file(): global PAINTING_ENTRANCES, PAINTING_EXITS + from . import datatypes + from Utils import safe_builtins + class RenameUnpickler(pickle.Unpickler): def find_class(self, module, name): - renamed_module = module - if module == "datatypes": - renamed_module = "worlds.lingo.datatypes" - - return super(RenameUnpickler, self).find_class(renamed_module, name) + if module in ("worlds.lingo.datatypes", "datatypes"): + return getattr(datatypes, name) + elif module == "builtins" and name in safe_builtins: + return getattr(safe_builtins, name) + raise pickle.UnpicklingError(f"global '{module}.{name}' is forbidden") file = pkgutil.get_data(__name__, os.path.join("data", "generated.dat")) pickdata = RenameUnpickler(BytesIO(file)).load() diff --git a/worlds/lingo/test/TestDatafile.py b/worlds/lingo/test/TestDatafile.py index 9f4e9da0..60acb3e8 100644 --- a/worlds/lingo/test/TestDatafile.py +++ b/worlds/lingo/test/TestDatafile.py @@ -1,8 +1,8 @@ import os import unittest -from worlds.lingo.static_logic import HASHES -from worlds.lingo.utils.pickle_static_data import hash_file +from ..static_logic import HASHES +from ..utils.pickle_static_data import hash_file class TestDatafile(unittest.TestCase):