Fixes after testing

tests/test_markup.py

@@ -3,13 +3,18 @@
 except ImportError:
     pass
 
+try:
+    import enum
+except ImportError:
+    enum = None
+
 from transitions.extensions.markup import MarkupMachine, rep
 from transitions.extensions.factory import HierarchicalMarkupMachine
 from .utils import Stuff
 from functools import partial
 
 
-from unittest import TestCase
+from unittest import TestCase, skipIf
 
 try:
     from unittest.mock import MagicMock
@@ -158,3 +163,23 @@ def setUp(self):
         self.machine_cls = HierarchicalMarkupMachine
         self.num_trans = len(self.transitions)
         self.num_auto = self.num_trans + 9**2
+
+
+@skipIf(enum is None, "enum is not available")
+class TestMarkupMachineEnum(TestMarkupMachine):
+    class States(enum.Enum):
+        A = 1
+        B = 2
+        C = 3
+        D = 4
+
+    def setUp(self):
+        self.machine_cls = MarkupMachine
+        self.states = TestMarkupMachineEnum.States
+        self.transitions = [
+            {'trigger': 'walk', 'source': self.states.A, 'dest': self.states.B},
+            {'trigger': 'run', 'source': self.states.B, 'dest': self.states.C},
+            {'trigger': 'sprint', 'source': self.states.C, 'dest': self.states.D}
+        ]
+        self.num_trans = len(self.transitions)
+        self.num_auto = self.num_trans + len(self.states)**2

transitions/extensions/diagrams.py

@@ -3,6 +3,7 @@
 
 import warnings
 import logging
+from enum import Enum
 from functools import partial
 
 _LOGGER = logging.getLogger(__name__)
@@ -164,7 +165,10 @@ def _get_graph(self, model, title=None, force_new=False, show_roi=False):
             grph = self.graph_cls(self, title=title if title is not None else self.title)
             self.model_graphs[model] = grph
             try:
-                self.model_graphs[model].set_node_style(model.state.name, 'active')
+                if isinstance(model.state, Enum):
+                    self.model_graphs[model].set_node_style(model.state.name, 'active')
+                else:
+                    self.model_graphs[model].set_node_style(model.state, 'active')
             except AttributeError:
                 _LOGGER.info("Could not set active state of diagram")
         try:

transitions/extensions/markup.py

@@ -1,8 +1,8 @@
-from enum import Enum
 from six import string_types, iteritems
 from functools import partial
 import itertools
 import importlib
+from enum import Enum
 from collections import defaultdict
 
 from ..core import Machine
@@ -68,7 +68,10 @@ def _convert_states(self, states):
         markup_states = []
         for state in states:
             s_def = _convert(state, self.state_attributes, self.skip_references)
-            s_def['name'] = state.name
+            if isinstance(state, Enum):
+                s_def['name'] = state.name
+            else:
+                s_def['name'] = getattr(state, '_name', state.name)
             if getattr(state, 'children', False):
                 s_def['children'] = self._convert_states(state.children)
             markup_states.append(s_def)