revng-revng/python/revng/ptml/parser.py

#
# This file is distributed under the MIT License. See LICENSE.md for details.
#
# noqa: A003

from __future__ import annotations

import xml.sax
from dataclasses import dataclass
from functools import cached_property
from typing import Dict, Generator, List

from revng.support import ToBytesInput, to_bytes


@dataclass
class Position:
    """Represents a single character in a multi-line text document"""

    line: int
    character: int

    def clone(self):
        return Position(self.line, self.character)

    def to_string(self):
        return f"{self.line}:{self.character}"


@dataclass
class Range:
    """Represents a range of characters in a multi-line text document. Both
    start and end are inclusive"""

    start: Position
    end: Position

    def is_single_line(self):
        return self.start.line == self.end.line

    def to_string(self):
        return f"{self.start.to_string()},{self.end.to_string()}"


class TextExtractor:
    """Utility function that allows extracting `Range`s from a text"""

    def __init__(self, text: str):
        # This is stored just so the lines property can be computed lazily via
        # the cached_property decorator
        self.original_text = text

    @cached_property
    def lines(self) -> List[str]:
        return self.original_text.split("\n")

    def get_text(self, range_: Range) -> str:
        if range_.is_single_line():
            return self._get_line(range_.start.line, range_.start.character, range_.end.character)
        lines = [self._get_line(range_.start.line, range_.start.character)]
        for line_index in range(range_.start.line, range_.end.line + 1):
            lines.append(self.lines[line_index])
        lines.append(self._get_line(range_.end.line, -1, range_.end.character))
        return "\n".join(lines)

    def _get_line(self, index: int, start=-1, end=-1):
        line = self.lines[index]
        if start == -1 and end == -1:
            return line
        elif start == -1:
            return line[:end]
        elif end == -1:
            return line[start:]
        else:
            return line[start:end]


@dataclass
class TreeNode:
    """The node of a tree of ranged attributes"""

    attributes: Dict[str, str]
    children: List[TreeNode]
    text_range: Range


@dataclass
class RangedAttributes:
    attributes: Dict[str, str]
    text_range: Range


class MetadataTree:
    """Utility class that allows interacting with a TreeNode tree more easily"""

    def __init__(self, root: TreeNode):
        self.root = root

    def walk(self) -> Generator[RangedAttributes, None, None]:
        queue = [self.root]
        index = 0
        while index < len(queue):
            top = queue[index]
            queue.extend(top.children)
            yield RangedAttributes(top.attributes, top.text_range)
            index += 1


class Parser(xml.sax.ContentHandler):
    """This ContentHandler will split the document into its bare text and a
    tree of TreeNodes stored in root"""

    def __init__(self):
        # The bare text
        self.text = ""
        # The root node of the document
        self.root = TreeNode({}, [], Range(Position(0, 0), Position(0, 0)))
        self._currentPosition = Position(0, 0)
        self._stack = [self.root]

    def startElement(self, name: str, attrs: xml.sax.xmlreader.AttributesImpl):  # noqa: N802
        assert name in ("div", "span")
        new_node = TreeNode(dict(attrs), [], Range(self._currentPosition.clone(), Position(0, 0)))
        self._stack[-1].children.append(new_node)
        self._stack.append(new_node)

    def endElement(self, name: str):  # noqa: N802
        # Mark the top-most node's end position
        self._stack[-1].text_range.end = self._currentPosition.clone()
        # Pop the top-most node from the stack
        self._stack.pop()

    def endDocument(self):  # noqa: N802
        # Check that we didn't miss closing an element
        assert len(self._stack) == 1 and self._stack[0] is self.root
        # Mark the end of the root node
        self.root.text_range.end = self._currentPosition.clone()

    def characters(self, content: str):
        self.text += content
        parts = content.split("\n")
        if len(parts) == 1:
            self._currentPosition.character += len(content)
        else:
            self._currentPosition.line += len(parts) - 1
            self._currentPosition.character = len(parts[-1])


@dataclass(frozen=True)
class PTMLDocument:
    text: str
    metadata: MetadataTree
    extractor: TextExtractor


def parse(input_: ToBytesInput) -> PTMLDocument:
    """Split a single PTML document into its text and metadata"""
    handler = Parser()
    with to_bytes(input_) as wrapped:
        xml.sax.parseString(wrapped, handler)
    return PTMLDocument(handler.text, MetadataTree(handler.root), TextExtractor(handler.text))