Source

yt_ajax / repl / repl.py

Full commit
import code, sys, StringIO, json
import tempfile
import threading
import webbrowser
from wsgiref.simple_server import make_server
import base64
from yt.utilities.amr_utils import write_png_to_file
from yt.visualization.image_writer import map_to_colors
import numpy as na

transport = []

class ManagedInterpreter(code.InteractiveInterpreter):
    def runsource(self, *args, **kwargs):
        old_stdout, old_stderr = sys.stdout, sys.stderr
        new_stdout = StringIO.StringIO()
        sys.stdout = new_stdout
        sys.stderr = new_stdout
        code.InteractiveInterpreter.runsource(self, *args, **kwargs)
        sys.stdout, sys.stderr = old_stdout, old_stderr
        new_stdout.seek(0)
        return new_stdout.read()

def add_image(buffer):
    image = na.log10(buffer)
    mi = na.nanmin(image[~na.isinf(image)])
    ma = na.nanmax(image[~na.isinf(image)])
    color_bounds = mi, ma
    image = (image - color_bounds[0])/(color_bounds[1] - color_bounds[0])
    to_plot = map_to_colors(image, "algae")
    to_plot = na.clip(to_plot, 0, 255)
    tf = tempfile.TemporaryFile()
    write_png_to_file(to_plot, tf)
    tf.seek(0)
    s = tf.read()
    tf.close()
    transport.append(s)

ma = ManagedInterpreter(locals=dict(add_image = add_image))

FILE = 'frontend.html'
PORT = 8080

def test_app(environ, start_response):
    if environ['REQUEST_METHOD'] == 'POST':
        try:
            request_body_size = int(environ['CONTENT_LENGTH'])
            request_body = environ['wsgi.input'].read(request_body_size).strip()
        except (TypeError, ValueError):
            request_body = "0"
        try:
            interp_resp = "%s\n" % (request_body)
            interp_resp += ma.runsource(request_body)
            interp_resp += ">>> "
            image = ""
            new_image = 0
            if len(transport) > 0:
                new_image = 1
                image = base64.encodestring(transport.pop(0))
            response_body = json.dumps(
                    dict(text=interp_resp, image=image, new_image = new_image)
                )
            #print response_body
        except KeyError:
            response_body = json.dumps(dict(text="error", image="", new_image=0))
        status = '200 OK'
        headers = [('Content-type', 'text/plain')]
        start_response(status, headers)
        return [response_body]
    else:
        response_body = open(FILE).read()
        status = '200 OK'
        headers = [('Content-type', 'text/html'),
                   ('Content-Length', str(len(response_body)))]
        start_response(status, headers)
        return [response_body]

def open_browser():
    """Start a browser after waiting for half a second."""
    def _open_browser():
        webbrowser.open('http://localhost:%s/%s' % (PORT, FILE))
    thread = threading.Timer(0.5, _open_browser)
    thread.start()

def start_server():
    """Start the server."""
    httpd = make_server("", PORT, test_app)
    httpd.serve_forever()

if __name__ == "__main__":
    open_browser()
    start_server()