Parse new colors for highlight and shader

.gitignore

@@ -1,2 +1,3 @@
 __pycache__/
 /.hypothesis/
+.DS_Store

src/rmscene/scene_items.py

@@ -76,8 +76,7 @@ class PenColor(enum.IntEnum):
     GRAY_OVERLAP = 8
 
     # All highlight colors share the same value.
-    # There is also yet unknown extra data in the block
-    # that might contain additional color information.
+    # This is a placeholder, see the colormap below for details.
     HIGHLIGHT = 9
 
     GREEN_2 = 10
@@ -86,6 +85,42 @@ class PenColor(enum.IntEnum):
 
     YELLOW_2 = 13
 
+    # HIGHLIGHT enumerated
+    HIGHLIGHT_YELLOW = 14
+    HIGHLIGHT_BLUE = 15
+    HIGHLIGHT_PINK = 16
+    HIGHLIGHT_ORANGE = 17
+    HIGHLIGHT_GREEN = 18
+    HIGHLIGHT_GRAY = 19
+
+    # SHADER enumerated
+    SHADER_GRAY = 20
+    SHADER_ORANGE = 21
+    SHADER_MAGENTA = 22
+    SHADER_BLUE = 23
+    SHADER_RED = 24
+    SHADER_GREEN = 25
+    SHADER_YELLOW = 26
+    SHADER_CYAN = 27
+
+# colors hardcoded in rm files for highlight and shader
+HARDCODED_COLORMAP = {
+    (255, 237, 117, 255): PenColor.HIGHLIGHT_YELLOW,
+    (190, 234, 254, 255): PenColor.HIGHLIGHT_BLUE,
+    (242, 158, 255, 255): PenColor.HIGHLIGHT_PINK,
+    (255, 195, 140, 255): PenColor.HIGHLIGHT_ORANGE,
+    (172, 255, 133, 255): PenColor.HIGHLIGHT_GREEN,
+    (199, 199, 198, 255): PenColor.HIGHLIGHT_GRAY,
+    (33, 30, 28, 64): PenColor.SHADER_GRAY,
+    (254, 178, 0, 115): PenColor.SHADER_ORANGE,
+    (192, 127, 210, 128): PenColor.SHADER_MAGENTA,
+    (48, 74, 224, 77): PenColor.SHADER_BLUE,
+    (194, 49, 50, 102): PenColor.SHADER_RED,
+    (145, 218, 113, 128): PenColor.SHADER_GREEN,
+    (250, 231, 25, 115): PenColor.SHADER_YELLOW,
+    (116, 210, 232, 102): PenColor.SHADER_CYAN,
+}
+
 
 @enum.unique
 class Pen(enum.IntEnum):

src/rmscene/scene_stream.py

@@ -6,12 +6,14 @@
 
 from __future__ import annotations
 
+import io
 import logging
 import math
 import typing as tp
 from abc import ABC, abstractmethod
 from collections.abc import Iterable, Iterator
 from dataclasses import KW_ONLY, dataclass, replace
+from typing import Optional
 from uuid import UUID, uuid4
 
 from packaging.version import Version
@@ -430,6 +432,22 @@ def line_from_stream(stream: TaggedBlockReader, version: int = 2) -> si.Line:
     else:
         move_id = None
 
+    # This is for the color information (highlight & shader)
+    if stream.bytes_remaining_in_block() >= 6:
+        # not sure what this is, seems fixed x84x01
+        unk = stream.data.read_bytes(2)
+        b = stream.data.read_uint8()
+        g = stream.data.read_uint8()
+        r = stream.data.read_uint8()
+        a = stream.data.read_uint8()
+        rgba = (r, g, b, a)
+
+        if unk != b"\x84\x01" or rgba not in si.HARDCODED_COLORMAP:
+            _logger.warning(f"Unhandled color {rgba} with prefix {unk}")
+            stream.data.data.seek(-6, io.SEEK_CUR)
+        else:
+            color = si.HARDCODED_COLORMAP[rgba]
+
     return si.Line(color, tool, points, thickness_scale, starting_length, move_id)
 
 
@@ -448,6 +466,14 @@ def line_to_stream(line: si.Line, writer: TaggedBlockWriter, version: int = 2):
     writer.write_id(6, timestamp)
     if line.move_id is not None:
         writer.write_id(7, line.move_id)
+
+    if line.color in si.HARDCODED_COLORMAP.values():
+        rgba = [key for key, value in si.HARDCODED_COLORMAP.items() if value == line.color][0]
+        writer.data.write_bytes(b"\x84\x01")
+        writer.data.write_uint8(rgba[2])
+        writer.data.write_uint8(rgba[1])
+        writer.data.write_uint8(rgba[0])
+        writer.data.write_uint8(rgba[3])
 
 
 @dataclass
@@ -492,6 +518,7 @@ def from_stream(cls, stream: TaggedBlockReader) -> SceneItemBlock:
                 assert item_type == subclass.ITEM_TYPE
                 value = subclass.value_from_stream(stream)
             # Keep known extra data from within the value subblock
+
             extra_value_data = block_info.extra_data
         else:
             value = None

tests/test_color_tool.py

@@ -1,35 +1,27 @@
-import logging
-import os
 from pathlib import Path
 
 import pytest
 
 from rmscene import SceneGlyphItemBlock, SceneLineItemBlock, read_blocks
 from rmscene.scene_items import Pen, PenColor
 
-logger = logging.getLogger(__name__)
-
-
 DATA_PATH = Path(__file__).parent / "data"
 
 
-def _hex_lines(b, n=32):
-    return [b[i * n : (i + 1) * n].hex() for i in range(len(b) // n + 1)]
-
-
-FILE_NAME = os.path.join(DATA_PATH, "Color_and_tool_v3.14.4.rm")
-
-
 @pytest.mark.parametrize(
     "block_type,colors,tools",
     [
         (SceneGlyphItemBlock, {PenColor.HIGHLIGHT}, None),
         (SceneLineItemBlock, {PenColor.HIGHLIGHT}, {Pen.SHADER}),
-        (SceneLineItemBlock, {PenColor.GREEN_2, PenColor.CYAN, PenColor.MAGENTA}, {Pen.BALLPOINT_2}),
+        (
+            SceneLineItemBlock,
+            {PenColor.GREEN_2, PenColor.CYAN, PenColor.MAGENTA},
+            {Pen.BALLPOINT_2},
+        ),
     ],
 )
 def test_color_tool_parsing(block_type, colors, tools):
-
+    FILE_NAME = DATA_PATH / "Color_and_tool_v3.14.4.rm"
     with open(FILE_NAME, "rb") as f:
         result = read_blocks(f)
         for el in result:
@@ -40,4 +32,43 @@ def test_color_tool_parsing(block_type, colors, tools):
             if isinstance(el, block_type) and el.item.value.color in colors:
                 if tools is None:
                     continue
-                assert el.item.value.tool in tools, "Tool and colors don't match"
+                assert el.item.value.tool in tools, "Tool and colors don't match"
+
+
+def test_highlight_shader_colors():
+    FILE_NAME = DATA_PATH / "More_color_highlight_shader_v3.15.4.2.rm"
+    with open(FILE_NAME, "rb") as f:
+        result = read_blocks(f)
+        expected_colors = [
+            PenColor.HIGHLIGHT_YELLOW,
+            PenColor.HIGHLIGHT_BLUE,
+            PenColor.HIGHLIGHT_PINK,
+            PenColor.HIGHLIGHT_ORANGE,
+            PenColor.HIGHLIGHT_GREEN,
+            PenColor.HIGHLIGHT_GRAY,
+            PenColor.SHADER_GRAY,
+            PenColor.SHADER_ORANGE,
+            PenColor.SHADER_MAGENTA,
+            PenColor.SHADER_BLUE,
+            PenColor.SHADER_RED,
+            PenColor.SHADER_GREEN,
+            PenColor.SHADER_YELLOW,
+            PenColor.SHADER_CYAN,
+            PenColor.BLACK,
+            PenColor.GRAY,
+            PenColor.WHITE,
+            PenColor.BLUE,
+            PenColor.RED,
+            PenColor.GREEN_2,
+            PenColor.YELLOW_2,
+            PenColor.CYAN,
+            PenColor.MAGENTA,
+        ]
+        start = 0
+        for block in result:
+            if isinstance(block, SceneLineItemBlock) and block.item.value:
+                assert (
+                    block.item.value.color == expected_colors[start]
+                ), f"Unexpected color {block.item.value.color} at index {start}"
+                start += 1
+        assert start == len(expected_colors)

tests/test_scene_stream.py

@@ -45,6 +45,7 @@ def _hex_lines(b, n=32):
     ("Wikipedia_highlighted_p2.rm", "3.1"),
     ("With_SceneInfo_Block.rm", "3.4"),  # XXX version?
     ("Color_and_tool_v3.14.4.rm", "3.14"),
+    ("More_color_highlight_shader_v3.15.4.2.rm", "3.15"),
 ]