From 9ceb319f9ff30fa6a2a393f142c765aa4a532acc Mon Sep 17 00:00:00 2001
From: Frederik Mathiesen <frederik_mathies@hotmail.com>
Date: Fri, 27 Mar 2020 00:50:34 -0700
Subject: [PATCH] PlantUML: add arrow from history to initial memory

---
 sismic/io/plantuml.py | 15 ++++++++++++++-
 1 file changed, 14 insertions(+), 1 deletion(-)

diff --git a/sismic/io/plantuml.py b/sismic/io/plantuml.py
index 1b41b33..6672cfb 100644
--- a/sismic/io/plantuml.py
+++ b/sismic/io/plantuml.py
@@ -2,7 +2,7 @@
 import re
 import sys
 
-from typing import Dict, List, Tuple, Optional
+from typing import Dict, List, Tuple, Union
 from ..io import import_from_yaml
 from ..model import (
     DeepHistoryState, FinalState, Transition, CompoundState,
@@ -170,6 +170,9 @@ def export_state(self, name: str) -> None:
 
         self.export_transitions(name)
 
+        if isinstance(state, (ShallowHistoryState, DeepHistoryState)):
+            self.export_history_memory(state)
+
         # Nested states
         for i, child in enumerate(self.statechart.children_for(name)):
             if i != 0 and isinstance(state, OrthogonalState):
@@ -226,6 +229,16 @@ def export_transition(self, transition: Transition) -> None:
             text=''.join(text),
         ))
 
+    def export_history_memory(self, history_state: Union[ShallowHistoryState, DeepHistoryState]):
+        if history_state.memory:
+            target = self.statechart.state_for(history_state.memory)
+
+            self.output('{source} {arrow} {target}'.format(
+                source=self.state_id(history_state.name),
+                arrow=self.arrow(history_state, target),
+                target=self.state_id(target.name)
+            ))
+
     def export(self) -> str:
         self.output('@startuml')