tmapdump.py

import enum
import io
import itertools
import json
import math
import numpy as np
import struct
import typing
import zlib

class Map:
    class Classification(enum.IntEnum):
        """Enumeration mapping the map pixel classifications."""
        UNKNOWN = 0b000
        TILE = 0b001
        WALL = 0b010
        WATER = 0b011
        LAVA = 0b100
        HONEY = 0b101
        FAR_BACKGROUND = 0b110
        NEAR_BACKGROUND = 0b111
    _elem_dtype = np.dtype([
        ('color_id', np.uint16),
        ('classification', np.uint8),
        ('paint_id', np.uint8),
        ('light', np.uint8),
    ])
    def __init__(self, f:typing.BinaryIO, world_surface=None, rock_layer=None):
        """Load a Map. world_surface and rock_layer are the corresponding Y positions, if known."""
        # Map header
        if struct.unpack('<4x7sB12x',f.read(24)) != (b'relogic', 1):
            raise RuntimeError("This doesn't appear to be a valid map file.")
        f.seek(f.read(1)[0],1) # Skip the world name
        height, width, n_tiles, n_walls, n_liquids, n_sky, n_dirt, n_rock = struct.unpack('<4x2I6H', f.read(24))
        def iter_bits(bytes_obj):
            for byte in bytes_obj:
                for n in (0x01, 0x02, 0x04, 0x08, 0x10, 0x20, 0x40, 0x80):
                    yield byte&n==n
        tile_has_multiple_options = f.read(math.ceil(n_tiles/8))
        wall_has_multiple_options = f.read(math.ceil(n_walls/8))
        tile_option_count = [(f.read(1)[0] if has_mo else 1) for i,has_mo in zip(range(n_tiles),iter_bits(tile_has_multiple_options))]
        wall_option_count = [(f.read(1)[0] if has_mo else 1) for i,has_mo in zip(range(n_walls),iter_bits(wall_has_multiple_options))]
        color_id_counter = itertools.count()
        unknown_id = next(color_id_counter)
        tile_id = tuple(itertools.chain.from_iterable([next(color_id_counter) for i in range(num_options)] for num_options in tile_option_count))
        wall_id = tuple(itertools.chain.from_iterable([next(color_id_counter) for i in range(num_options)] for num_options in wall_option_count))
        water_id = next(color_id_counter)
        lava_id = next(color_id_counter)
        honey_id = next(color_id_counter)
        sky_id = tuple(next(color_id_counter) for i in range(n_sky))
        dirt_id = tuple(next(color_id_counter) for i in range(n_dirt))
        rock_id = tuple(next(color_id_counter) for i in range(n_rock))
        hell_id = next(color_id_counter)
        n_color_ids = next(color_id_counter)

        # Map data
        if world_surface is None:
            world_surface = height*0.3 # TODO Calculate based on rock_layer if present
        if rock_layer is None:
            rock_layer = world_surface + height*0.2
        def tile_iter(f):
            x = 0
            y = 0
            read_u16 = struct.Struct('<H')
            while True:
                flags = f.read(1)[0]
                has_paint = flags&0x01==0x01
                classification = (flags&0x0E)>>1
                word_type = flags&0x10==0x10
                light_saved = flags&0x20==0x20
                rle_byte = flags&0x40==0x40
                rle_word = flags&0x80==0x80

                if has_paint:
                    paint_id = f.read(1)[0] >> 1
                else:
                    paint_id = 0

                if classification in {0b001, 0b010, 0b111}:
                    if word_type:
                        tile_type, = read_u16.unpack(f.read(2))
                    else:
                        tile_type = f.read(1)[0]
                else:
                    tile_type = 0

                if light_saved:
                    light = f.read(1)[0]
                elif classification==0b000:
                    light = 0
                else:
                    light = 255

                if classification==0b000:
                    color_id = unknown_id
                elif classification==0b001:
                    color_id = tile_id[tile_type]
                elif classification==0b010:
                    color_id = wall_id[tile_type]
                elif classification==0b011:
                    color_id = water_id
                elif classification==0b100:
                    color_id = lava_id
                elif classification==0b101:
                    color_id = honey_id
                elif classification==0b110:
                    if y<world_surface:
                        color_id = sky_id[math.floor(n_sky * (y / world_surface))]
                    else:
                        color_id = hell_id
                elif classification==0b111:
                    if y<rock_layer:
                        color_id = dirt_id[tile_type]
                    else:
                        color_id = rock_id[tile_type]

                yield (color_id, classification, paint_id, light)
                x += 1

                if rle_byte or rle_word:
                    if rle_byte:
                        rle_count = f.read(1)[0]
                    else:
                        rle_count, = read_u16.unpack(f.read(2))
                    if light_saved:
                        for light in f.read(rle_count):
                            yield (color_id, classification, paint_id, light)
                    else:
                        for i in range(rle_count):
                            yield (color_id, classification, paint_id, light)
                    x += rle_count
                y += x//width
                x -= (x//width)*width
        with io.BytesIO(zlib.decompress(f.read(),wbits=-15)) as uncompressed_stream:
            self._data = np.fromiter(tile_iter(uncompressed_stream), dtype=Map._elem_dtype, count=width*height)
            assert self._data.flags.c_contiguous
            self._data.shape = (height, width)
            assert self._data.flags.c_contiguous
            self._data.flags.writeable = False
    @property
    def color_ids(self)->np.ndarray:
        """Get an array of color IDs."""
        return self._data['color_id']
    @property
    def classifications(self)->np.ndarray:
        """Get an array of classifications."""
        return self._data['classification']
    @property
    def paint_ids(self)->np.ndarray:
        """Get an array of paint IDs."""
        return self._data['paint_id']
    @property
    def light_levels(self)->np.ndarray:
        """Get an array of light levels."""
        return self._data['light']

class GameInfo:
    """Class responsible for holding information from Terraria's information dump."""

    class Construct:
        """Base class holding general information on tiles and walls from the game."""
        _name: str
        _properties: typing.FrozenSet[str]
        _color_ids: typing.Tuple[int, ...]
        def __init__(self, json_obj:dict):
            self._name = json_obj['Name']
            self._properties = frozenset(json_obj['Sets'])
            self._color_ids = tuple(json_obj['ColorIDs'])
        @property
        def name(self)->str:
            """Internal name for the construct."""
            return self._name
        @property
        def properties(self)->typing.FrozenSet[str]:
            """Collection of various attributes the construct has."""
            return self._properties
        @property
        def color_ids(self)->typing.Tuple[int, ...]:
            """Palette indexes for each variation of the construct."""
            return self._color_ids
        def __repr__(self)->str:
            return f"<{type(self).__name__}: {self._name}>"
    class Tile(Construct):
        """Class holding information about a tile from the game."""
        pass
    class Wall(Construct):
        """Class holding information about a wall from the game."""
        pass

    class Cell:
        """Base class that represents a map cell's block type, excluding paint and light."""
        def __eq__(self, other)->bool:
            return isinstance(other, type(self))
        def __ne__(self, other)->bool:
            return not self==other
        def __hash__(self)->int:
            return 0
        def __repr__(self)->str:
            return f"<{type(self).__name__}>"
    class VariantCell(Cell):
        """Base class for a cell holding a block that exists in variations."""
        _variant: int
        def __init__(self, variant:int):
            self._variant = variant
        @property
        def variant(self)->int:
            """The particular variation held by the cell."""
            return self._variant
        def __eq__(self, other)->bool:
            return super().__eq__(other) and self._variant==other._variant
        def __hash__(self)->int:
            return self._variant
        def __repr__(self)->str:
            return f"<{type(self).__name__}: #{self._variant}>"
    class UnknownCell(Cell):
        """Cell that doesn't have any information known, other than it being within bounds of the map."""
        pass
    class TileCell(VariantCell):
        """Cell that contains a tile."""
        _tile: 'GameInfo.Tile'
        def __init__(self, variant:int, tile:'GameInfo.Tile'):
            super().__init__(variant)
            self._tile = tile
        @property
        def tile(self)->'GameInfo.Tile':
            """Information about the contained tile."""
            return self._tile
        def __eq__(self, other)->bool:
            return super().__eq__(other) and self._tile==other._tile
        def __hash__(self)->int:
            return hash((self._tile,self.variant))
        def __repr__(self)->str:
            return f"<{type(self).__name__}: {self._tile.name} #{self._variant}>"
    class WallCell(VariantCell):
        """Cell that contains a wall."""
        _wall: 'GameInfo.Wall'
        def __init__(self, variant:int, wall:'GameInfo.Wall'):
            super().__init__(variant)
            self._wall = wall
        @property
        def wall(self)->'GameInfo.Wall':
            """Information about the contained wall."""
            return self._wall
        def __eq__(self, other)->bool:
            return super().__eq__(other) and self._wall==other._wall
        def __hash__(self)->int:
            return hash((self._wall,self.variant))
        def __repr__(self)->str:
            return f"<{type(self).__name__}: {self._wall.name} #{self._variant}>"
    class LiquidCell(Cell):
        """Base class for cells containing liquid."""
        pass
    class WaterCell(LiquidCell):
        """Cell that contains water."""
        pass
    class LavaCell(LiquidCell):
        """Cell that contains lava."""
        pass
    class HoneyCell(LiquidCell):
        """Cell that contains honey."""
        pass
    class BackgroundCell(Cell):
        """Base class for cells that represent the background."""
        pass
    class SkyCell(VariantCell, BackgroundCell):
        """Cell that shows the sky."""
        pass
    class DirtCell(VariantCell, BackgroundCell):
        """Cell that shows a dirt background."""
        pass
    class RockCell(VariantCell, BackgroundCell):
        """Cell that shows a rock background."""
        pass
    class HellCell(BackgroundCell):
        """Cell that shows hell."""
        pass

    _color_id_to_cell: typing.Dict[int, Cell]
    _default_palette: np.ndarray
    _tiles: typing.Tuple[Tile, ...]
    _walls: typing.Tuple[Wall, ...]
    _paint_names: typing.Tuple[str, ...]
    _default_paint_palette: np.ndarray
    def __init__(self, f:typing.BinaryIO):
        json_object = json.load(f)

        self._color_id_to_cell = {}
        self._color_id_to_cell[json_object['EmptyColorID']] = GameInfo.UnknownCell()
        self._tiles = tuple(map(GameInfo.Tile, json_object['Tiles']))
        for tile_obj in self._tiles:
            for variant, color_id in zip(itertools.count(), tile_obj.color_ids):
                self._color_id_to_cell[color_id] = GameInfo.TileCell(variant, tile_obj)
        self._walls = tuple(map(GameInfo.Wall, json_object['Walls']))
        for wall_obj in self._walls:
            for variant, color_id in zip(itertools.count(), wall_obj.color_ids):
                self._color_id_to_cell[color_id] = GameInfo.WallCell(variant, wall_obj)
        for cell, color_id in zip(
                (GameInfo.WaterCell(), GameInfo.LavaCell(), GameInfo.HoneyCell()),
                json_object['LiquidColorIDs']):
            self._color_id_to_cell[color_id] = cell
        for variant, color_id in zip(itertools.count(), json_object['SkyColorIDs']):
            self._color_id_to_cell[color_id] = GameInfo.SkyCell(variant)
        for variant, color_id in zip(itertools.count(), json_object['DirtColorIDs']):
            self._color_id_to_cell[color_id] = GameInfo.DirtCell(variant)
        for variant, color_id in zip(itertools.count(), json_object['RockColorIDs']):
            self._color_id_to_cell[color_id] = GameInfo.RockCell(variant)
        self._color_id_to_cell[json_object['HellColorID']] = GameInfo.HellCell()

        self._default_palette = np.array(
                [[x['R'],x['G'],x['B'],x['A']] for x in json_object['Colors']],
                dtype=np.uint8)
        self._paint_names = tuple(x['Name'] for x in json_object['Paint'])
        self._default_paint_palette = np.array(
                [[x['Color']['R'],x['Color']['G'],x['Color']['B'],x['Color']['A']] for x in json_object['Paint']],
                dtype=np.uint8)
    @property
    def default_palette(self)->np.ndarray:
        """Color palette for cells as extracted from the game."""
        return self._default_palette
    def make_palette(self, fn: typing.Callable[['GameInfo.Cell'],np.ndarray])->np.ndarray:
        """Make a color palette; fn takes a Cell and returns the "color" to use for that position."""
        palette = [None]*len(self._color_id_to_cell)
        for color_id, cell in self._color_id_to_cell.items():
            palette[color_id] = np.array(fn(cell))
        return np.array(palette)
    @property
    def tiles(self)->typing.Tuple[Tile, ...]:
        """Sequence of Tiles in index order."""
        return self._tiles
    @property
    def walls(self)->typing.Tuple[Wall, ...]:
        """Sequence of Walls in index order."""
        return self._walls
    @property
    def paint_names(self)->typing.Tuple[str, ...]:
        """Sequence of paint internal names in index order."""
        return self._paint_names
    @property
    def default_paint_palette(self)->np.ndarray:
        """Paint color palette as extracted from the game. Does not account for special effects."""
        return self._default_paint_palette
    
    def render_as_game(self, map_obj:Map)->np.ndarray:
        """Render map_obj as the game probably would."""
        # Based on Terraria.Map.MapHelper.GetMapTileXnaColor in 1.4.2
        base_color = self.default_palette[map_obj.color_ids].transpose([2,0,1]) # Move RGBA to the outer dimension
        paint_id = map_obj.paint_ids
        base_color_rgb = base_color[0:-1]
        # BEGIN Terraria.Map.MapHelper.MapColor
        paint_color = self.default_paint_palette[paint_id].transpose([2,0,1]) # Move RGBA again
        paint_color_rgb = paint_color[0:-1]
        x, y, z = base_color_rgb
        x = np.fmax(x,y)
        x,z = np.fmax(x,z), np.fmin(x,z)
        mapcolor_case_31_rgb = base_color_rgb
        mapcolor_case_29_rgb = np.floor(paint_color_rgb * (z * 0.3))
        mapcolor_case_30_wall_rgb = np.floor((255-base_color_rgb)*0.5)
        mapcolor_case_30_tile_rgb = 255-base_color_rgb
        mapcolor_case_30_rgb = np.where(
                map_obj.classifications==Map.Classification.WALL,
                mapcolor_case_30_wall_rgb,
                mapcolor_case_30_tile_rgb)
        mapcolor_rgb = np.floor(paint_color_rgb * x)
        mapcolor_rgb = np.where(paint_id==31,mapcolor_case_31_rgb,mapcolor_rgb)
        mapcolor_rgb = np.where(paint_id==29,mapcolor_case_29_rgb,mapcolor_rgb)
        mapcolor_rgb = np.where(paint_id==30,mapcolor_case_30_rgb,mapcolor_rgb)
        # END Terraria.Map.MapHelper.MapColor
        base_color_rgb = np.where(paint_id==0,base_color_rgb,mapcolor_rgb)
        dark_color_rgb = np.floor(base_color_rgb * (map_obj.light_levels/255.0))
        use_base = (map_obj.light_levels==255) | (paint_id==31)
        color_rgb = np.where(use_base,base_color_rgb,dark_color_rgb)
        color_rgba = np.concatenate([color_rgb,base_color[-1,np.newaxis]],axis=0)
        return color_rgba.transpose([1,2,0]) # Move RGBA back to the outer layer

    def color_id_to_cell(self, color_id:int)->Cell:
        """Convert a color ID to a Cell"""
        return self._color_id_to_cell[color_id]