From d40c7d9f09d9b012837d6060a4c598b23b19646f Mon Sep 17 00:00:00 2001
From: Martin Ashby <martin@martin-laptop.lan>
Date: Thu, 17 May 2018 18:18:11 +0100
Subject: Palette, tools, load & save working

---
 server.py | 50 ++++++++++++++++++++++++++++++++++++++++++++------
 1 file changed, 44 insertions(+), 6 deletions(-)

(limited to 'server.py')

diff --git a/server.py b/server.py
index ea4b814..e610f57 100644
--- a/server.py
+++ b/server.py
@@ -3,6 +3,9 @@ from flask import Flask, send_from_directory
 from flask_sockets import Sockets
 from gevent import pywsgi
 from geventwebsocket.handler import WebSocketHandler
+import os
+import os.path
+import pickle
 
 try:
     import unicornhathd as unicorn
@@ -10,24 +13,31 @@ try:
 except ImportError:
     from unicorn_hat_sim import unicornhathd as unicorn
 
+# Actions that can be sent 
 NO_OP = 'NO_OP'
 SET_PIXEL = 'SET_PIXEL'
 CLEAR = 'CLEAR'
+SAVE = 'SAVE'
+LOAD = 'LOAD'
 
+# GLobal references to app & clients 
 all_clients = set()
-
 app = Flask(__name__, static_folder='build/')
 sockets = Sockets(app)
 
 def send_state():
     global all_clients
-    jsonPixels = json.dumps(unicorn.get_pixels())
+    state = {
+        "saves": os.listdir(SAVES_DIR),
+        "pixels": unicorn.get_pixels()
+    }
+    stateJson = json.dumps(state)
     for client in all_clients:
-        client.send(jsonPixels)
+        client.send(stateJson)
 
 def execute_command(command):
     cmd_type = command['type']
-    if cmd_type == NO_OP:
+    if cmd_type == NO_OP: # Used just to get state without doing an action
         pass 
     elif cmd_type == SET_PIXEL:
         x = int(command['x'])
@@ -40,7 +50,28 @@ def execute_command(command):
     elif cmd_type == CLEAR:
         unicorn.clear()
         unicorn.show()
+    elif cmd_type == SAVE:
+        saveName = command["saveName"]
+        save(saveName)
+    elif cmd_type == LOAD:
+        saveName = command["saveName"]
+        load(saveName)
+
+# Constants
+SAVES_DIR = 'saves/'
 
+def save(saveName):
+    with open(os.path.join(SAVES_DIR, saveName), "wb") as f:
+        pickle.dump(unicorn.get_pixels(), f)
+
+def load(saveName):
+    with open(os.path.join(SAVES_DIR, saveName), "rb") as f:
+        pixels = pickle.load(f)
+        for x, row in enumerate(pixels):
+            for y, pixel in enumerate(row):
+                unicorn.set_pixel(x, y, *pixel)
+        unicorn.show()
+    
 @sockets.route('/ws')
 def do_websocket(websocket):
     global all_clients
@@ -48,6 +79,8 @@ def do_websocket(websocket):
     try:
         while not websocket.closed:
             command = websocket.receive()
+            if not command: 
+                break
             cmd = json.loads(command)
             execute_command(cmd)
             send_state()
@@ -60,12 +93,17 @@ def send_static(path):
     return send_from_directory('build/', path)
 
 def main():
+    try:
+        os.mkdir(SAVES_DIR)
+    except:
+        pass
     try:
         print("Serving on port 3001")
         server = pywsgi.WSGIServer(('', 3001), app, handler_class=WebSocketHandler)
         server.serve_forever()
-    except:
+    finally:
         unicorn.off()
 
 if __name__=="__main__":
-    main()
\ No newline at end of file
+    main()
+    
\ No newline at end of file
-- 
cgit v1.2.3-ZIG