python-pages 0.1.0__py3-none-any.whl

pages/tables.py

@@ -0,0 +1,493 @@
+"""Middle-layer reader and conservative writer for Pages TST tables."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, fields
+import struct
+from typing import TYPE_CHECKING
+import uuid
+
+from .components import GraphClone, clone_object_graph, replace_uuid_identities
+from .protobuf import Message
+from .stylesheet import CharacterProperties, character_style_properties
+from .text import LocatedObject, parse_attribute_table, parse_reference
+
+if TYPE_CHECKING:
+    from .pages import PagesDocument
+
+
+TYPE_TABLE_INFO = 6000
+TYPE_TABLE_MODEL = 6001
+TYPE_TILE = 6002
+TYPE_TABLE_DATA_LIST = 6005
+TYPE_ATTACHMENT = 2003
+TYPE_HEADER_STORAGE_BUCKET = 6006
+TYPE_COLUMN_ROW_UID_MAP = 6267
+TYPE_STROKE_SIDECAR = 6305
+
+STORAGE_TABLE_ATTACHMENTS = 9
+
+TABLE_INFO_MODEL = 2
+TABLE_MODEL_DATA_STORE = 4
+TABLE_MODEL_ROWS = 6
+TABLE_MODEL_COLUMNS = 7
+TABLE_MODEL_NAME = 8
+TABLE_MODEL_HEADER_ROWS = 9
+TABLE_MODEL_HEADER_COLUMNS = 10
+TABLE_MODEL_FOOTER_ROWS = 11
+TABLE_MODEL_BODY_TEXT_STYLE = 24
+TABLE_MODEL_HEADER_ROW_TEXT_STYLE = 25
+TABLE_MODEL_HEADER_COLUMN_TEXT_STYLE = 26
+TABLE_MODEL_FOOTER_ROW_TEXT_STYLE = 27
+TABLE_MODEL_COLUMN_ROW_UIDS = 46
+TABLE_MODEL_STROKE_SIDECAR = 49
+
+DATA_STORE_ROW_HEADERS = 1
+DATA_STORE_COLUMN_HEADERS = 2
+DATA_STORE_TILES = 3
+DATA_STORE_STRING_TABLE = 4
+DATA_STORE_STYLE_TABLE = 5
+
+DATA_LIST_TYPE = 1
+DATA_LIST_NEXT_KEY = 2
+DATA_LIST_ENTRIES = 3
+DATA_LIST_STRING = 1
+DATA_LIST_STYLE = 4
+LIST_ENTRY_KEY = 1
+LIST_ENTRY_REFCOUNT = 2
+LIST_ENTRY_STRING = 3
+LIST_ENTRY_REFERENCE = 4
+
+CELL_STORAGE_VERSION = 5
+CELL_TYPE_TEXT = 3
+
+
+@dataclass(frozen=True)
+class CellStorage:
+    """The portions of a version-5 TST cell buffer used by this API."""
+
+    raw: bytes
+    cell_type: int
+    string_id: int | None = None
+    rich_text_id: int | None = None
+    cell_style_id: int | None = None
+    text_style_id: int | None = None
+
+
+def _reference(message: Message, number: int) -> int | None:
+    field = message.one(number)
+    return None if field is None else parse_reference(bytes(field.value))
+
+
+def _parse_cell_storage(data: bytes) -> CellStorage:
+    if len(data) < 12:
+        raise ValueError("TST cell storage is shorter than its 12-byte header")
+    if data[0] != CELL_STORAGE_VERSION:
+        raise ValueError(f"unsupported TST cell storage version {data[0]}")
+    flags = struct.unpack_from("<I", data, 8)[0]
+    offset = 12
+    identifiers: dict[int, int] = {}
+    # Every active flag contributes one fixed-width payload in increasing flag
+    # order. Numeric/date values precede the keyed identifiers.
+    widths = {
+        0x1: 16,
+        0x2: 8,
+        0x4: 8,
+        **{1 << bit: 4 for bit in range(3, 21)},
+    }
+    for flag, width in widths.items():
+        if not flags & flag:
+            continue
+        if offset + width > len(data):
+            raise ValueError("truncated TST cell storage flag payload")
+        if width == 4:
+            identifiers[flag] = struct.unpack_from("<i", data, offset)[0]
+        offset += width
+    return CellStorage(
+        raw=data,
+        cell_type=data[1],
+        string_id=identifiers.get(0x8),
+        rich_text_id=identifiers.get(0x10),
+        cell_style_id=identifiers.get(0x20),
+        text_style_id=identifiers.get(0x40),
+    )
+
+
+def _row_buffers(row_info: Message, columns: int) -> list[bytes | None]:
+    storage = row_info.one(6)
+    offsets_field = row_info.one(7)
+    if storage is None or offsets_field is None:
+        raise ValueError("pre-BNC TST table rows are unsupported")
+    offsets_raw = bytes(offsets_field.value)
+    if len(offsets_raw) % 2:
+        raise ValueError("TST cell offset array has an odd byte length")
+    offsets = list(struct.unpack(f"<{len(offsets_raw) // 2}h", offsets_raw))
+    wide = row_info.one(8)
+    if wide is not None and bool(wide.value):
+        offsets = [value * 4 for value in offsets]
+    payload = bytes(storage.value)
+    result: list[bytes | None] = []
+    for column in range(columns):
+        if column >= len(offsets) or offsets[column] < 0:
+            result.append(None)
+            continue
+        start = offsets[column]
+        end = next(
+            (value for value in offsets[column + 1 :] if value >= 0),
+            len(payload),
+        )
+        result.append(payload[start:end])
+    return result
+
+
+class TableStore:
+    """Resolved TableInfo/TableModel/DataStore graph for one body table."""
+
+    def __init__(self, document: "PagesDocument", info: LocatedObject):
+        self.document = document
+        self.info = info
+        model_id = _reference(info.message, TABLE_INFO_MODEL)
+        model = document.objects().get(model_id) if model_id is not None else None
+        if model is None or model.info.type_id != TYPE_TABLE_MODEL:
+            raise ValueError(f"TST.TableInfoArchive {info.segment.identifier} has no table model")
+        self.model = model
+        data_store = model.message.one(TABLE_MODEL_DATA_STORE)
+        if data_store is None:
+            raise ValueError(f"TST.TableModelArchive {model_id} has no DataStore")
+        self.data_store = Message.parse(bytes(data_store.value))
+        self.rows = int(model.message.one(TABLE_MODEL_ROWS).value)
+        self.columns = int(model.message.one(TABLE_MODEL_COLUMNS).value)
+        name = model.message.one(TABLE_MODEL_NAME)
+        self.name = bytes(name.value).decode("utf-8", "replace") if name else ""
+        self._cells = self._load_cells()
+
+    def _load_cells(self) -> dict[tuple[int, int], CellStorage]:
+        tiles_field = self.data_store.one(DATA_STORE_TILES)
+        if tiles_field is None:
+            return {}
+        tiles = Message.parse(bytes(tiles_field.value))
+        tile_size_field = tiles.one(2)
+        tile_size = int(tile_size_field.value) if tile_size_field is not None else 256
+        result: dict[tuple[int, int], CellStorage] = {}
+        objects = self.document.objects()
+        for tile_field in tiles.get(1):
+            tile_link = Message.parse(bytes(tile_field.value))
+            tile_id_field = tile_link.one(1)
+            tile_id = int(tile_id_field.value) if tile_id_field is not None else 0
+            object_id = _reference(tile_link, 2)
+            tile = objects.get(object_id) if object_id is not None else None
+            if tile is None or tile.info.type_id != TYPE_TILE:
+                continue
+            for row_field in tile.message.get(5):
+                row_info = Message.parse(bytes(row_field.value))
+                row_index_field = row_info.one(1)
+                if row_index_field is None:
+                    continue
+                row = tile_id * tile_size + int(row_index_field.value)
+                for column, raw in enumerate(_row_buffers(row_info, self.columns)):
+                    if raw is not None:
+                        result[(row, column)] = _parse_cell_storage(raw)
+        return result
+
+    def cell_storage(self, row: int, column: int) -> CellStorage | None:
+        return self._cells.get((row, column))
+
+    def _data_list(self, field_number: int, expected_type: int) -> LocatedObject:
+        object_id = _reference(self.data_store, field_number)
+        obj = self.document.objects().get(object_id) if object_id is not None else None
+        if obj is None or obj.info.type_id != TYPE_TABLE_DATA_LIST:
+            raise ValueError(f"table DataStore field {field_number} has no TableDataList")
+        list_type = obj.message.one(DATA_LIST_TYPE)
+        if list_type is None or int(list_type.value) != expected_type:
+            raise ValueError(f"TableDataList {object_id} has unexpected list type")
+        return obj
+
+    @staticmethod
+    def _entry(data_list: LocatedObject, key: int) -> tuple[object, Message]:
+        for raw_entry in data_list.message.get(DATA_LIST_ENTRIES):
+            entry = Message.parse(bytes(raw_entry.value))
+            key_field = entry.one(LIST_ENTRY_KEY)
+            if key_field is not None and int(key_field.value) == key:
+                return raw_entry, entry
+        raise KeyError(f"TableDataList {data_list.segment.identifier} has no key {key}")
+
+    def string(self, key: int) -> str:
+        data_list = self._data_list(DATA_STORE_STRING_TABLE, DATA_LIST_STRING)
+        _, entry = self._entry(data_list, key)
+        value = entry.one(LIST_ENTRY_STRING)
+        return "" if value is None else bytes(value.value).decode("utf-8", "replace")
+
+    def set_string(self, key: int, value: str) -> None:
+        data_list = self._data_list(DATA_STORE_STRING_TABLE, DATA_LIST_STRING)
+        raw_entry, entry = self._entry(data_list, key)
+        refcount = entry.one(LIST_ENTRY_REFCOUNT)
+        if refcount is not None and int(refcount.value) != 1:
+            raise NotImplementedError("shared TST string entries require cell-key reassignment")
+        string = entry.one(LIST_ENTRY_STRING)
+        if string is None:
+            entry.add_bytes(LIST_ENTRY_STRING, value.encode("utf-8"))
+        else:
+            string.set_bytes(value.encode("utf-8"))
+        raw_entry.set_bytes(entry.encode())
+
+    def text_style(self, row: int, column: int, cell: CellStorage | None) -> LocatedObject | None:
+        style_key = cell.text_style_id if cell is not None else None
+        if style_key is not None:
+            styles = self._data_list(DATA_STORE_STYLE_TABLE, DATA_LIST_STYLE)
+            _, entry = self._entry(styles, style_key)
+            object_id = _reference(entry, LIST_ENTRY_REFERENCE)
+        else:
+            header_rows = self.model.message.one(TABLE_MODEL_HEADER_ROWS)
+            header_columns = self.model.message.one(TABLE_MODEL_HEADER_COLUMNS)
+            footer_rows = self.model.message.one(TABLE_MODEL_FOOTER_ROWS)
+            if header_rows is not None and row < int(header_rows.value):
+                field_number = TABLE_MODEL_HEADER_ROW_TEXT_STYLE
+            elif header_columns is not None and column < int(header_columns.value):
+                field_number = TABLE_MODEL_HEADER_COLUMN_TEXT_STYLE
+            elif footer_rows is not None and row >= self.rows - int(footer_rows.value):
+                field_number = TABLE_MODEL_FOOTER_ROW_TEXT_STYLE
+            else:
+                field_number = TABLE_MODEL_BODY_TEXT_STYLE
+            object_id = _reference(self.model.message, field_number)
+        return self.document.objects().get(object_id) if object_id is not None else None
+
+    def character_properties(
+        self, row: int, column: int, cell: CellStorage | None
+    ) -> CharacterProperties:
+        style = self.text_style(row, column, cell)
+        chain: list[CharacterProperties] = []
+        seen: set[int] = set()
+        while style is not None and style.segment.identifier not in seen:
+            seen.add(style.segment.identifier)
+            chain.append(character_style_properties(style))
+            super_field = style.message.one(1)
+            if super_field is None:
+                break
+            super_message = Message.parse(bytes(super_field.value))
+            parent_id = _reference(super_message, 3)
+            style = self.document.objects().get(parent_id) if parent_id is not None else None
+        values = {}
+        for item in fields(CharacterProperties):
+            value = next(
+                (getattr(props, item.name) for props in chain if getattr(props, item.name) is not None),
+                None,
+            )
+            values[item.name] = value
+        return CharacterProperties(**values)
+
+
+def body_table_stores(document: "PagesDocument", storage: LocatedObject) -> list[TableStore]:
+    """Resolve body attachment runs to their TST table graphs in document order."""
+    field = storage.message.one(STORAGE_TABLE_ATTACHMENTS)
+    if field is None:
+        return []
+    objects = document.objects()
+    result: list[TableStore] = []
+    seen: set[int] = set()
+    for run in parse_attribute_table(bytes(field.value)):
+        attachment = objects.get(run.object_id) if run.object_id is not None else None
+        if attachment is None or attachment.info.type_id != TYPE_ATTACHMENT:
+            continue
+        info_id = _reference(attachment.message, 1)
+        info = objects.get(info_id) if info_id is not None else None
+        if info is None or info.info.type_id != TYPE_TABLE_INFO or info_id in seen:
+            continue
+        seen.add(info_id)
+        result.append(TableStore(document, info))
+    return result
+
+
+def _body_table_attachment(
+    document: "PagesDocument", storage: LocatedObject
+) -> LocatedObject:
+    field = storage.message.one(STORAGE_TABLE_ATTACHMENTS)
+    objects = document.objects()
+    if field is not None:
+        for run in parse_attribute_table(bytes(field.value)):
+            attachment = objects.get(run.object_id) if run.object_id is not None else None
+            if attachment is None or attachment.info.type_id != TYPE_ATTACHMENT:
+                continue
+            info_id = _reference(attachment.message, 1)
+            info = objects.get(info_id) if info_id is not None else None
+            if info is not None and info.info.type_id == TYPE_TABLE_INFO:
+                return attachment
+    raise ValueError("document has no Pages-authored table graph to use as a template")
+
+
+def _set_varint(message: Message, number: int, value: int) -> None:
+    field = message.one(number)
+    if field is None:
+        message.add_varint(number, value)
+    else:
+        field.set_varint(value)
+
+
+def _resize_header_bucket(
+    bucket: Message, count: int, cells_per_header: int
+) -> None:
+    source = bucket.get(2)
+    if not source:
+        return
+    template = Message.parse(bytes(source[0].value))
+    bucket.fields = [field for field in bucket.fields if field.number != 2]
+    for index in range(count):
+        header = Message.parse(template.encode())
+        _set_varint(header, 1, index)
+        _set_varint(header, 4, cells_per_header)
+        bucket.add_bytes(2, header.encode())
+
+
+def _uuid_message(value: uuid.UUID) -> bytes:
+    message = Message()
+    message.add_varint(1, value.int >> 64)
+    message.add_varint(2, value.int & ((1 << 64) - 1))
+    return message.encode()
+
+
+def _rebuild_uid_map(
+    uid_map: Message, rows: int, columns: int, *, seed: str
+) -> None:
+    """Replace the cloned 2x2 row/column UID index with the requested grid."""
+
+    def dimension(count: int, label: str) -> tuple[list[bytes], list[int], list[int]]:
+        indexed = [
+            (index, uuid.uuid5(uuid.NAMESPACE_URL, f"{seed}:{label}:{index}"))
+            for index in range(count)
+        ]
+        sorted_values = sorted(indexed, key=lambda item: item[1].int)
+        index_for_uid = [index for index, _ in sorted_values]
+        uid_for_index = [0] * count
+        for uid_index, (original_index, _) in enumerate(sorted_values):
+            uid_for_index[original_index] = uid_index
+        return (
+            [_uuid_message(value) for _, value in sorted_values],
+            index_for_uid,
+            uid_for_index,
+        )
+
+    columns_sorted, column_indexes, column_uids = dimension(columns, "column")
+    rows_sorted, row_indexes, row_uids = dimension(rows, "row")
+    uid_map.fields = []
+    for encoded in columns_sorted:
+        uid_map.add_bytes(1, encoded)
+    for value in column_indexes:
+        uid_map.add_varint(2, value)
+    for value in column_uids:
+        uid_map.add_varint(3, value)
+    for encoded in rows_sorted:
+        uid_map.add_bytes(4, encoded)
+    for value in row_indexes:
+        uid_map.add_varint(5, value)
+    for value in row_uids:
+        uid_map.add_varint(6, value)
+
+
+def _clear_tile(tile: Message, rows: int, columns: int) -> None:
+    tile.fields = [field for field in tile.fields if field.number != 5]
+    for number, value in ((1, 0), (2, 0), (3, 0), (4, rows), (6, 5), (7, 1), (8, 1)):
+        _set_varint(tile, number, value)
+    offsets = struct.pack(f"<{columns}h", *([-1] * columns))
+    for row_index in range(rows):
+        row = Message()
+        row.add_varint(1, row_index)
+        row.add_varint(2, 0)
+        row.add_bytes(3, "🤠".encode("utf-8"))
+        row.add_bytes(4, "🤠".encode("utf-8"))
+        row.add_varint(5, 5)
+        row.add_bytes(6, b"")
+        row.add_bytes(7, offsets)
+        row.add_varint(8, 1)
+        tile.add_bytes(5, row.encode())
+
+
+def clone_body_table(
+    document: "PagesDocument", storage: LocatedObject, rows: int, columns: int
+) -> tuple[TableStore, GraphClone]:
+    """Clone the fixture-authored complete TST graph as a blank body table."""
+    template_attachment = _body_table_attachment(document, storage)
+
+    def include(obj: LocatedObject) -> bool:
+        return obj.filename.startswith("Index/CalculationEngine-") or obj.filename.startswith(
+            "Index/Tables/"
+        )
+
+    clone = clone_object_graph(
+        document,
+        template_attachment.segment.identifier,
+        include=include,
+    )
+    objects = document.objects()
+    info_id = _reference(clone.root.message, 1)
+    info = objects.get(info_id) if info_id is not None else None
+    if info is None or info.info.type_id != TYPE_TABLE_INFO:
+        raise ValueError("cloned table attachment has no TableInfoArchive")
+    model_id = _reference(info.message, TABLE_INFO_MODEL)
+    model = objects.get(model_id) if model_id is not None else None
+    if model is None or model.info.type_id != TYPE_TABLE_MODEL:
+        raise ValueError("cloned TableInfoArchive has no TableModelArchive")
+
+    seed = f"python-pages:table:{info.segment.identifier}:{model.segment.identifier}"
+    for new_identifier in clone.identifiers.values():
+        replace_uuid_identities(objects[new_identifier].message, seed=seed)
+    table_id = model.message.one(1)
+    generated_table_id = str(uuid.uuid5(uuid.NAMESPACE_URL, seed)).upper().encode("ascii")
+    if table_id is None:
+        model.message.add_bytes(1, generated_table_id)
+    else:
+        table_id.set_bytes(generated_table_id)
+    _set_varint(model.message, TABLE_MODEL_ROWS, rows)
+    _set_varint(model.message, TABLE_MODEL_COLUMNS, columns)
+    name = model.message.one(TABLE_MODEL_NAME)
+    table_number = len(body_table_stores(document, storage)) + 1
+    encoded_name = f"Table {table_number}".encode("utf-8")
+    if name is None:
+        model.message.add_bytes(TABLE_MODEL_NAME, encoded_name)
+    else:
+        name.set_bytes(encoded_name)
+
+    data_store_field = model.message.one(TABLE_MODEL_DATA_STORE)
+    if data_store_field is None:
+        raise ValueError("cloned TableModelArchive has no DataStore")
+    data_store = Message.parse(bytes(data_store_field.value))
+    row_headers_field = data_store.one(DATA_STORE_ROW_HEADERS)
+    if row_headers_field is None:
+        raise ValueError("cloned table DataStore has no row HeaderStorage")
+    row_headers = Message.parse(bytes(row_headers_field.value))
+    row_bucket_ids = [
+        parse_reference(bytes(field.value)) for field in row_headers.get(2)
+    ]
+    column_bucket_id = _reference(data_store, DATA_STORE_COLUMN_HEADERS)
+    for row_bucket_id in row_bucket_ids:
+        bucket = objects.get(row_bucket_id) if row_bucket_id is not None else None
+        if bucket is None or bucket.info.type_id != TYPE_HEADER_STORAGE_BUCKET:
+            raise ValueError("cloned table row HeaderStorageBucket is missing")
+        _resize_header_bucket(bucket.message, rows, columns)
+    column_bucket = (
+        objects.get(column_bucket_id) if column_bucket_id is not None else None
+    )
+    if column_bucket is None or column_bucket.info.type_id != TYPE_HEADER_STORAGE_BUCKET:
+        raise ValueError("cloned table column HeaderStorageBucket is missing")
+    _resize_header_bucket(column_bucket.message, columns, rows)
+
+    uid_map_id = _reference(model.message, TABLE_MODEL_COLUMN_ROW_UIDS)
+    uid_map = objects.get(uid_map_id) if uid_map_id is not None else None
+    if uid_map is None or uid_map.info.type_id != TYPE_COLUMN_ROW_UID_MAP:
+        raise ValueError("cloned TableModelArchive has no ColumnRowUIDMapArchive")
+    _rebuild_uid_map(uid_map.message, rows, columns, seed=seed)
+
+    sidecar_id = _reference(model.message, TABLE_MODEL_STROKE_SIDECAR)
+    sidecar = objects.get(sidecar_id) if sidecar_id is not None else None
+    if sidecar is None or sidecar.info.type_id != TYPE_STROKE_SIDECAR:
+        raise ValueError("cloned TableModelArchive has no StrokeSidecarArchive")
+    _set_varint(sidecar.message, 2, columns)
+    _set_varint(sidecar.message, 3, rows)
+
+    # Every owned tile is made blank. Retaining the Pages-authored list/style
+    # components is intentional: the empty tile no longer references their
+    # entries, but the complete editing graph remains available to Pages.
+    for old, new in clone.identifiers.items():
+        obj = objects[new]
+        if obj.info.type_id == TYPE_TILE:
+            _clear_tile(obj.message, rows, columns)
+
+    return TableStore(document, info), clone