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