Update chat demo to be more exemplary.

bdarnell · bdarnell · commit 974c229c8bf5 · 2013-03-16T00:52:04.000-04:00
Use gen.coroutine in auth handler (but not in wait_for_messages
because it makes things less clear and we don't have good cancellation
support).  No more mixin.  Python 3 compatible.
diff --git a/demos/chat/chatdemo.py b/demos/chat/chatdemo.py
@@ -18,107 +18,93 @@
 import tornado.auth
 import tornado.escape
 import tornado.ioloop
-import tornado.options
 import tornado.web
 import os.path
 import uuid
 
-from tornado.options import define, options
+from tornado import gen
+from tornado.options import define, options, parse_command_line
 
 define("port", default=8888, help="run on the given port", type=int)
 
 
-class Application(tornado.web.Application):
+class MessageBuffer(object):
     def __init__(self):
-        handlers = [
-            (r"/", MainHandler),
-            (r"/auth/login", AuthLoginHandler),
-            (r"/auth/logout", AuthLogoutHandler),
-            (r"/a/message/new", MessageNewHandler),
-            (r"/a/message/updates", MessageUpdatesHandler),
-        ]
-        settings = dict(
-            cookie_secret="__TODO:_GENERATE_YOUR_OWN_RANDOM_VALUE_HERE__",
-            login_url="/auth/login",
-            template_path=os.path.join(os.path.dirname(__file__), "templates"),
-            static_path=os.path.join(os.path.dirname(__file__), "static"),
-            xsrf_cookies=True,
-        )
-        tornado.web.Application.__init__(self, handlers, **settings)
-
-
-class BaseHandler(tornado.web.RequestHandler):
-    def get_current_user(self):
-        user_json = self.get_secure_cookie("chatdemo_user")
-        if not user_json: return None
-        return tornado.escape.json_decode(user_json)
-
-
-class MainHandler(BaseHandler):
-    @tornado.web.authenticated
-    def get(self):
-        self.render("index.html", messages=MessageMixin.cache)
-
-
-class MessageMixin(object):
-    waiters = set()
-    cache = []
-    cache_size = 200
+        self.waiters = set()
+        self.cache = []
+        self.cache_size = 200
 
     def wait_for_messages(self, callback, cursor=None):
-        cls = MessageMixin
         if cursor:
-            index = 0
-            for i in xrange(len(cls.cache)):
-                index = len(cls.cache) - i - 1
-                if cls.cache[index]["id"] == cursor: break
-            recent = cls.cache[index + 1:]
-            if recent:
-                callback(recent)
+            new_count = 0
+            for msg in reversed(self.cache):
+                if msg["id"] == cursor:
+                    break
+                new_count += 1
+            if new_count:
+                callback(self.cache[-new_count:])
                 return
-        cls.waiters.add(callback)
+        self.waiters.add(callback)
 
     def cancel_wait(self, callback):
-        cls = MessageMixin
-        cls.waiters.remove(callback)
+        self.waiters.remove(callback)
 
     def new_messages(self, messages):
-        cls = MessageMixin
-        logging.info("Sending new message to %r listeners", len(cls.waiters))
-        for callback in cls.waiters:
+        logging.info("Sending new message to %r listeners", len(self.waiters))
+        for callback in self.waiters:
             try:
                 callback(messages)
             except:
                 logging.error("Error in waiter callback", exc_info=True)
-        cls.waiters = set()
-        cls.cache.extend(messages)
-        if len(cls.cache) > self.cache_size:
-            cls.cache = cls.cache[-self.cache_size:]
+        self.waiters = set()
+        self.cache.extend(messages)
+        if len(self.cache) > self.cache_size:
+            self.cache = self.cache[-self.cache_size:]
+
 
+# Making this a non-singleton is left as an exercise for the reader.
+global_message_buffer = MessageBuffer()
 
-class MessageNewHandler(BaseHandler, MessageMixin):
+
+class BaseHandler(tornado.web.RequestHandler):
+    def get_current_user(self):
+        user_json = self.get_secure_cookie("chatdemo_user")
+        if not user_json: return None
+        return tornado.escape.json_decode(user_json)
+
+
+class MainHandler(BaseHandler):
+    @tornado.web.authenticated
+    def get(self):
+        self.render("index.html", messages=global_message_buffer.cache)
+
+
+class MessageNewHandler(BaseHandler):
     @tornado.web.authenticated
     def post(self):
         message = {
             "id": str(uuid.uuid4()),
             "from": self.current_user["first_name"],
             "body": self.get_argument("body"),
         }
-        message["html"] = self.render_string("message.html", message=message)
+        # to_basestring is necessary for Python 3's json encoder,
+        # which doesn't accept byte strings.
+        message["html"] = tornado.escape.to_basestring(
+            self.render_string("message.html", message=message))
         if self.get_argument("next", None):
             self.redirect(self.get_argument("next"))
         else:
             self.write(message)
-        self.new_messages([message])
+        global_message_buffer.new_messages([message])
 
 
-class MessageUpdatesHandler(BaseHandler, MessageMixin):
+class MessageUpdatesHandler(BaseHandler):
     @tornado.web.authenticated
     @tornado.web.asynchronous
     def post(self):
         cursor = self.get_argument("cursor", None)
-        self.wait_for_messages(self.on_new_messages,
-                               cursor=cursor)
+        global_message_buffer.wait_for_messages(self.on_new_messages,
+                                                cursor=cursor)
 
     def on_new_messages(self, messages):
         # Closed client connection
@@ -127,23 +113,21 @@ def on_new_messages(self, messages):
         self.finish(dict(messages=messages))
 
     def on_connection_close(self):
-        self.cancel_wait(self.on_new_messages)
+        global_message_buffer.cancel_wait(self.on_new_messages)
 
 
 class AuthLoginHandler(BaseHandler, tornado.auth.GoogleMixin):
     @tornado.web.asynchronous
+    @gen.coroutine
     def get(self):
         if self.get_argument("openid.mode", None):
-            self.get_authenticated_user(self.async_callback(self._on_auth))
+            user = yield self.get_authenticated_user()
+            self.set_secure_cookie("chatdemo_user",
+                                   tornado.escape.json_encode(user))
+            self.redirect("/")
             return
         self.authenticate_redirect(ax_attrs=["name"])
 
-    def _on_auth(self, user):
-        if not user:
-            raise tornado.web.HTTPError(500, "Google auth failed")
-        self.set_secure_cookie("chatdemo_user", tornado.escape.json_encode(user))
-        self.redirect("/")
-
 
 class AuthLogoutHandler(BaseHandler):
     def get(self):
@@ -152,8 +136,21 @@ def get(self):
 
 
 def main():
-    tornado.options.parse_command_line()
-    app = Application()
+    parse_command_line()
+    app = tornado.web.Application(
+        [
+            (r"/", MainHandler),
+            (r"/auth/login", AuthLoginHandler),
+            (r"/auth/logout", AuthLogoutHandler),
+            (r"/a/message/new", MessageNewHandler),
+            (r"/a/message/updates", MessageUpdatesHandler),
+            ],
+        cookie_secret="__TODO:_GENERATE_YOUR_OWN_RANDOM_VALUE_HERE__",
+        login_url="/auth/login",
+        template_path=os.path.join(os.path.dirname(__file__), "templates"),
+        static_path=os.path.join(os.path.dirname(__file__), "static"),
+        xsrf_cookies=True,
+        )
     app.listen(options.port)
     tornado.ioloop.IOLoop.instance().start()
 
