Running code snippets in a background thread

Question

I'm working on a dataset from a MOOC. I have a lot of python3 code snippets that I need to run and get the results from. To do this I've written a python script that loops over each snippet. For each snippet I:

1. Create new StringIO objects
2. Set sys.stdout and sys.stderr to my stringIO buffers
3. Execute the code snippet in a threading.thread object
4. Join the thread
5. Log the results in the stringIO buffers
6. Restore stdout and stderr

This works fine for "correct" code, but this has undesired side effects in other cases:

• When the code has an infinite loop, thread.join doesn't kill the thread. The thread itself is a daemon thread, so it runs quietly in the background until my loop finishes.
• When the code has an infinite loop with a print(), the thread starts overwriting my actual stdout when I set it back to the default (away from the StringIO buffer). This pollutes my reporting.

Here is my current code:

def execCode(code, testScript=None):
    # create file-like string to capture output
    codeOut = io.StringIO()
    codeErr = io.StringIO()

    # capture output and errors
    sys.stdout = codeOut
    sys.stderr = codeErr

    def worker():
        exec(code, globals())

        if testScript:
            # flush stdout/stderror
            sys.stdout.truncate(0)
            sys.stdout.seek(0)
            # sys.stderr.truncate(0)
            # sys.stderr.seek(0)
            exec(testScript)

    thread = threading.Thread(target=worker, daemon=True)
    # thread = Process(target=worker) #, stdout=codeOut, stderr=codeErr)
    thread.start()
    thread.join(0.5)  # 500ms

    execError = codeErr.getvalue().strip()
    execOutput = codeOut.getvalue().strip()

    if thread.is_alive():
        thread.terminate()
        execError = "TimeError: run time exceeded"

    codeOut.close()
    codeErr.close()

    # restore stdout and stderr
    sys.stdout = sys.__stdout__
    sys.stderr = sys.__stderr__

    # restore any overridden functions
    restoreBuiltinFunctions()

    if execError:
        return False, stripOuterException(execError)
    else:
        return True, execOutput

To handle the undesirable cases, I've tried to use multithreading.Process and/or contextlib.redirect_stdout to run the code in a process (then I can call process.terminate()), but I'm not having any success capturing stdout/stderr.

So my question is: What can I do to make this better/more robust to handle bad code snippets?

(And yes, I know this is a bad idea in general; I'm running it in a virtual machine just in case there is malicious code in there somewhere)

Python version is 3.5.3

---

Update

It occurs to me that there is a little more flexibility in this situation. I have a function, preprocess(code) that accepts a the code submission as a string and alters it. Mostly I've been using it to swap out the value of some variables using regular expressions.

Here is an example implementation:

def preprocess(code):
    import re
    rx = re.compile('earlier_date\s*=\s*.+')
    code = re.sub(rx, "earlier_date = date(2016, 5, 3)", code)
    rx = re.compile('later_date\s*=\s*.+')
    code = re.sub(rx, "later_date = date(2016, 5, 24)", code)
    return code

I could use the preprocess function to help redirect STDOUT

Answer 1

Sorry, didn't have time to look at the issue of preventing the infinite loops, but here is a way to address the printing issue.

Enforce thread-local printing

As you've found, because each thread shares sys.stdout there isn't a good way to redirect printing from distinct threads to distinct streams. As you discovered, and as is documented, contextlib.redirect_stdout does not work in threaded applications, and has no effect on the output of subprocesses. However, because you're using Python 3 print is a function, and because Python is a dynamic language, we can (ab)use that.

Consider the following function signature:

def print(*objects, sep='', end='\n', file=sys.stdout, flush=False)

An astute reader may notice that this is Python 3's print function. Suppose we changed it just a little bit, to not use sys.stdout as the default stream and instead use a thread-local stream. Now we may be tempted to use the builtin threading.local function, however that will make it hard to recover the output after the fact. Let's start by making a mechanism for getting a thread-local output stream.

thread_local_data = {
    'stdout': {},
    'stderr': {}
}

def register_thread(thread_ident=None):
    """Register a thread with the monitor."""

    thread_ident = thread_ident or threading.get_ident()

    thread_local_data['stdout'][thread_ident] = io.StringIO()
    thread_local_data['stderr'][thread_ident] = io.StringIO()

def unregister_thread(thread_ident=None):
    """Unregister a thread with the monitor."""

    thread_ident = thread_ident or threading.get_ident()

    thread_local_data['stdout'][thread_ident].close()
    thread_local_data['stderr'][thread_ident].close()

def get_thread_stdout(thread_ident=None):
    """Get output stream for a thread."""

    thread_ident = thread_ident or threading.get_ident()
    return thread_local_data['stdout'][thread_ident]

def get_thread_stderr(thread_ident=None):
    """Get error stream for a thread."""

    thread_ident = thread_ident or threading.get_ident()
    return thread_local_data['stderr'][thread_ident]

original_print = print
def print(*objects, sep='', end='\n', file=None, flush=False):
    if file is None or file is sys.stdout:
        file = get_thread_stdout()
    if file is sys.stderr:
        file = get_thread_stderr()

    original_print(*objects, sep=sep, end=end, file=file, flush=flush)

You can test a simplified version of it right from the REPL - notice that we now have a print function in our globals() dict, where there was not one previously. Then we can see

>>> old_print=globals()['__builtins__'].print
>>> old_print("test")
test
>>> def print(): old_print("I WIN")
...
>>> import pprint
>>> pprint.pprint(globals())
{'__annotations__': {},
 '__builtins__': <module 'builtins' (built-in)>,
 '__doc__': None,
 '__loader__': <class '_frozen_importlib.BuiltinImporter'>,
 '__name__': '__main__',
 '__package__': None,
 '__spec__': None,
 'old_print': <built-in function print>,
 'pprint': <module 'pprint' from 'C:\\Users\\dobermil\\AppData\\Local\\Programs\\Python\\Python36-32\\lib\\pprint.py'>,
 'print': <function print at 0x02BD5198>}
>>> exec("print()", globals())
I WIN

Now that we have the mechanism to print as we want to, you just need to make sure you register the thread before starting it. That will setup the thread-local stdout.