1 files changed, 30 insertions, 7 deletions
diff --git a/examples/statemachine/trafficLightDemo.py b/examples/statemachine/trafficLightDemo.py
index 30fe934..ea42180 100644
--- a/examples/statemachine/trafficLightDemo.py
+++ b/examples/statemachine/trafficLightDemo.py
@@ -1,12 +1,35 @@
+#
+# trafficLightDemo.py
+#
+# Example of a simple state machine modeling the state of a traffic light
+#
+
 import statemachine
-import trafficlight
+import trafficlightstate
+
+
+class TrafficLight:
+    def __init__(self):
+        # start light in Red state
+        self._state = trafficlightstate.Red()
+
+    def change(self):
+        self._state = self._state.next_state()
+
+    # get light behavior/properties from current state
+    def __getattr__(self, attrname):
+        return getattr(self._state, attrname)
+
+    def __str__(self):
+        return "{}: {}".format(self.__class__.__name__, self._state)
+
 
-tl = trafficLight.Red()
+light = TrafficLight()
 for i in range(10):
-    print(tl, end='')
-    print(("STOP", "GO")[tl.carsCanGo])
-    tl.crossingSignal()
-    tl.delay()
+    print(light, end=' ')
+    print(("STOP", "GO")[light.carsCanGo])
+    light.crossingSignal()
+    light.delay()
     print()
 
-    tl = tl.nextState()
+    light.change()