-
Notifications
You must be signed in to change notification settings - Fork 1
/
Copy pathmain.py
171 lines (145 loc) · 7.07 KB
/
main.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
#
# Copyright 2019 Genentech Inc.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
"""
Genentech Warp Watcher
Minimal Tornado App to watch a directory and update with new files.
Author: Benjamin Barad <[email protected]>/<[email protected]>
"""
import datetime
import json
import os
import sys
from tornado.httpserver import HTTPServer
import tornado.ioloop
import tornado.template
from tornado.web import Application, RequestHandler, StaticFileHandler
from tornado.websocket import WebSocketHandler
from tornado.options import options, define, parse_command_line
define('port', default=8182, help='Port to listen on')
define('thumbnail_count', default=200, help="Number of thumbnails to show")
define('reload_rate', default=120000, help="How often to reload thumbnails, in ms")
define('parent_path', type=str, help="Parent directory for warp jobs")
# Set up globals for use throughout the program
static_path = {"path": "/tmp/thumbnails"}
warp_path = ["/tmp"]
body_html = [f"<h4>No Micrographs Loaded Yet - Current Warp Path is {warp_path[0]}</h4>"]
settings = {"websocket_ping_interval": 30}
clients = set()
loader = tornado.template.Loader(".")
string_loader = loader.load("template.html")
options.parse_command_line()
async def change_warp_path(new_directory):
"""
Change to a different warp output folder
Args:
new_directory(str): directory name for new warp folder to watch
Returns:
bool: true if successful change of folder.
"""
if not os.path.exists(os.path.join(options.parent_path, new_directory, "thumbnails")):
print("No Thumbnail Directory Found - Aborting")
return False
else:
warp_path[0] = os.path.join(options.parent_path, new_directory)
static_path["path"] = os.path.join(options.parent_path, new_directory, "thumbnails")
print(f"New Warp Directory: {warp_path[0]}")
return True
async def update_html_string(path, linecount=200):
"""
Generate HTML for the gallery of last n classes. Stores in a global variable to minimize
Args:
path (str): Path to warp folder
linecount (int): Number of thumbnails to display
"""
thumbnail_dir = os.path.join(path, "thumbnails")
if not os.path.exists(thumbnail_dir):
await message_all_clients({"type": "alert", "data": "Could not find thumbnails in the chosen directory"})
return
files = [file for file in os.listdir(thumbnail_dir) if (file.lower().endswith('.png'))]
files.sort(key=lambda x: os.path.getctime(os.path.join(thumbnail_dir, x)), reverse=True)
total_count = len(files)
if linecount > total_count:
linecount = total_count
return_list = [{"number": len(files)-i,"timestamp": datetime.datetime.fromtimestamp(os.path.getctime(os.path.join(thumbnail_dir, j))).strftime('%Y-%m-%d %H:%M:%S'), "name": j[:-4], "url": os.path.join("/gallery/", j)} for i,j in enumerate(files[:linecount])]
warp_name = os.path.split(path)[1]
if len(files) == 0:
body_html[0] = f"<h4>No Thumbnails Loaded Yet - Current Warp Path is {path}</h4>"
else:
body_html[0] = string_loader.generate(thumbnail_list = return_list, warp_name = warp_name, line_count = linecount, total_count = total_count).decode("utf-8")
print("Updating Client HTML")
await message_all_clients({"type": "gallery", "data": body_html[0]})
async def message_all_clients(message, clients = clients):
"""
Send a message to all open clients, and close any that respond with closed state.
Args:
message (str or dict): a websocket-friendly message.
clients (set): a set of websockethandler instances that should represent every open session.
"""
for client in list(clients):
try:
client.write_message(message)
except WebSocketClosedError:
logging.warn(f"Could not write a message to client {client} due to a WebSocketClosedError. Removing that client from the client list.")
clients.remove(client)
class IndexHandler(RequestHandler):
"""Core class to respond to new clients. Extends :py:class:`tornado.web.RequestHandler`"""
def get(self):
"""Minimal handler for setting up the very first connection via an HTTP request before setting up the websocket connection for all future interactions."""
self.render("index.html")
class SocketHandler(WebSocketHandler):
"""Primary Web Server control class - every new client will make initialize of these classes.
Extends :py:class:`tornado.websocket.WebSocketHandler`
"""
def open(self):
"""Adds new client to a global clients set when socket is opened."""
clients.add(self)
print("Socket Opened from {}".format(self.request.remote_ip))
async def on_message(self, message):
"""Handles messages from clients
Args:
message (str): string encoded json from the client. Should be of the format ``{"command": $COMMAND, "data": $DATA}``
"""
message_json = json.loads(message)
command = message_json['command']
data = message_json['data']
if command == "change_directory":
if data == None:
return
if await change_warp_path(data):
print(warp_path[0])
await update_html_string(warp_path[0], linecount=options.thumbnail_count)
else:
self.write_message({"type": "alert", "data": "Couldn't find a thumbnails folder in that directory - has Warp run there?"})
elif command == "initialize":
print("initializing with body html:")
self.write_message({"type": "gallery", "data": body_html[0]})
# await update_html_string(warp_path[0], linecount=options.thumbnail_count)
else:
self.write_message({"type": "alert", "data": "Didn't understand that message"})
def on_close(self):
"""Remove sockets from the clients list to minimize errors."""
clients.remove(self)
print("Socket Closed from {}".format(self.request.remote_ip))
def main():
"""Construct and serve the tornado app"""
app = Application([(r"/", IndexHandler), (r"/gallery/(.*)", StaticFileHandler, static_path), (r"/static/(.*)", StaticFileHandler, {"path": os.path.join(sys.path[0], "static")}), (r"/websocket", SocketHandler)], **settings)
app.listen(options.port)
listening_callback = tornado.ioloop.PeriodicCallback(lambda: update_html_string(warp_path[0], linecount=options.thumbnail_count), options.reload_rate)
listening_callback.start()
print(f'Listening on http://localhost:{options.port}')
tornado.ioloop.IOLoop.current().start()
if __name__ == "__main__":
main()