Miki Tebeka avatar Miki Tebeka committed 0c33475

start of redirect

Comments (0)

Files changed (1)

 __author__ = 'Miki Tebeka <miki.tebeka@gmail.com>'
 __version__ = '0.2.0'
 
-from six import b, exec_
+from contextlib import contextmanager
+from six import b, exec_, string_types
 from six.moves import socketserver as socksrv
 from threading import Thread
 from traceback import format_exc
 import socket
+import sys
 
 EOF = chr(4)
 DEFAULT_PORT = 9998
 
+@contextmanager
+def redirect_stdout(fo):
+    old_stdout = sys.stdout
+    sys.stdout = fo
+    yield fo
+    sys.stdout = old_stdout
+
+
+class Writer(object):
+    def __init__(self, wfile):
+        self.wfile = wfile
+
+    def write(self, obj):
+        if isinstance(obj, string_types):
+            obj = b(obj)
+        self.wfile.write(obj)
+
 
 class PyHandler(socksrv.StreamRequestHandler):
     password = None
     env = {}
+    redirect = False
 
     def handle(self):
         env = self.env.copy()
                 expr = self.rfile.readline().rstrip()
                 if expr == EOF:
                     return
-                try:
-                    value = eval(expr, globals(), env)
-                    self.wfile.write(format(value) + '\n')
-                except:
-                    exec_(expr, env)
+
+                out = Writer(self.wfile) if self.redirect else sys.stdout
+                with redirect_stdout(out):
+                    try:
+                        value = eval(expr, globals(), env)
+                        self.wfile.write(format(value) + '\n')
+                        self.wfile.flush()
+                    except:
+                        exec_(expr, env)
             except (EOFError, SystemExit, socket.error):
                 return
             except Exception:
                 error = format_exc()
                 self.wfile.write(b(error))
 
+
     def finish(self):
         try:
             socksrv.StreamRequestHandler.finish(self)
     allow_reuse_address = True
 
 
-def server_thread(env, port, password=None):
+def server_thread(env, port, password=None, redirect=False):
     # Create a new handler class for this with it's own env
     class Handler(PyHandler):
         pass
     Handler.env = env
     Handler.password = password
+    Handler.redirect = redirect
 
     server = ThreadedServer(('localhost', port), Handler)
     server.serve_forever()
 
 
-def install(env=None, port=DEFAULT_PORT, password=None):
+def install(env=None, port=DEFAULT_PORT, password=None, redirect=False):
     env = env or {}
-    t = Thread(target=server_thread, args=(env, port, password))
+    t = Thread(target=server_thread, args=(env, port, password, redirect))
     t.daemon = True
     t.start()
 
 
 
 def main(argv=None):
-    import sys
     from argparse import ArgumentParser
 
     argv = argv or sys.argv
                         default=DEFAULT_PORT)
     parser.add_argument('-l', '--login', help='login password',
                         default=None)
+    parser.add_argument('-r', '--redirect', help='redirect stdout',
+                        default=False, action='store_true')
     args = parser.parse_args(argv[1:])
 
-    t = install(port=args.port, password=args.login)
+    t = install(port=args.port, password=args.login, redirect=args.redirect)
     print('Serving on port {0}'.format(args.port))
     t.join()
 
Tip: Filter by directory path e.g. /media app.js to search for public/media/app.js.
Tip: Use camelCasing e.g. ProjME to search for ProjectModifiedEvent.java.
Tip: Filter by extension type e.g. /repo .js to search for all .js files in the /repo directory.
Tip: Separate your search with spaces e.g. /ssh pom.xml to search for src/ssh/pom.xml.
Tip: Use ↑ and ↓ arrow keys to navigate and return to view the file.
Tip: You can also navigate files with Ctrl+j (next) and Ctrl+k (previous) and view the file with Ctrl+o.
Tip: You can also navigate files with Alt+j (next) and Alt+k (previous) and view the file with Alt+o.