1from typing import List, Optional, Sequence
2
3from draftjs_exporter.command import Command
4from draftjs_exporter.constants import ENTITY_TYPES
5from draftjs_exporter.dom import DOM
6from draftjs_exporter.error import ExporterException
7from draftjs_exporter.options import Options, OptionsMap
8from draftjs_exporter.types import (
9    Block,
10    Element,
11    EntityDetails,
12    EntityKey,
13    EntityMap,
14)
15
16
17class EntityException(ExporterException):
18    pass
19
20
21class EntityState(object):
22    __slots__ = (
23        "entity_options",
24        "entity_map",
25        "entity_stack",
26        "completed_entity",
27        "element_stack",
28    )
29
30    def __init__(
31        self, entity_options: OptionsMap, entity_map: EntityMap
32    ) -> None:
33        self.entity_options = entity_options
34        self.entity_map = entity_map
35
36        self.entity_stack: List[EntityKey] = []
37        self.completed_entity: Optional[EntityKey] = None
38        self.element_stack: List[Element] = []
39
40    def apply(self, command: Command) -> None:
41        if command.name == "start_entity":
42            self.entity_stack.append(command.data)
43        elif command.name == "stop_entity":
44            expected_entity = self.entity_stack[-1]
45
46            if command.data != expected_entity:
47                raise EntityException(
48                    f"Expected {expected_entity}, got {command.data}"
49                )
50
51            self.completed_entity = self.entity_stack.pop()
52
53    def has_entity(self) -> List[EntityKey]:
54        return self.entity_stack
55
56    def has_no_entity(self) -> bool:
57        return not self.entity_stack
58
59    def get_entity_details(self, entity_key: EntityKey) -> EntityDetails:
60        details = self.entity_map.get(entity_key)
61
62        if details is None:
63            raise EntityException(
64                f'Entity "{entity_key}" does not exist in the entityMap'
65            )
66
67        return details
68
69    def render_entities(
70        self, style_node: Element, block: Block, blocks: Sequence[Block]
71    ) -> Element:
72        # We have a complete (start, stop) entity to render.
73        if self.completed_entity is not None:
74            entity_details = self.get_entity_details(self.completed_entity)
75            options = Options.get(
76                self.entity_options,
77                entity_details["type"],
78                ENTITY_TYPES.FALLBACK,
79            )
80            props = entity_details["data"].copy()
81            props["entity"] = {
82                "type": entity_details["type"],
83                "mutability": entity_details["mutability"]
84                if "mutability" in entity_details
85                else None,
86                "block": block,
87                "blocks": blocks,
88                "entity_range": {"key": self.completed_entity},
89            }
90
91            if len(self.element_stack) == 1:
92                children = self.element_stack[0]
93            else:
94                children = DOM.create_element()
95
96                for n in self.element_stack:
97                    DOM.append_child(children, n)
98
99            self.completed_entity = None
100            self.element_stack = []
101
102            # Is there still another entity? (adjacent) if so add the current style_node for it.
103            if self.has_entity():
104                self.element_stack.append(style_node)
105
106            return DOM.create_element(options.element, props, children)
107
108        if self.has_entity():
109            self.element_stack.append(style_node)
110            return None
111
112        return style_node
113