opdb/io_wrapper.py

89 lines
1.8 KiB
Python

import sys
import os
is_python3 = sys.version_info[0] == 3
input_queue = []
class RawInputException(Exception):
pass
class StreamStdin(object):
def __init__(self, input_queue):
self.input_queue = input_queue
def readline(self):
if len(self.input_queue) > 0:
return self.input_queue.pop(0)
else:
return ""
def raw_input_wrapper(*args):
if len(args) > 0:
print(args[0])
if input_queue:
return input_queue.pop(0)
raise RawInputException()
def python2_input_wrapper():
# Python 2 input() does eval(raw_input())
if input_queue:
input_str = input_queue.pop(0)
return eval(input_str)
raise RawInputException()
def stream_stdin():
while input_queue:
item = input_queue.pop(0)
yield item
def override_input():
if is_python3:
import builtins
else:
import __builtin__ as builtins
with open("user_code/input.txt") as file:
global input_queue
input_queue = list(file.readlines())
input_queue_new = []
for val in input_queue:
input_queue_new.append(val.replace("\n", "").strip())
input_queue = input_queue_new
builtins.raw_input = raw_input_wrapper
if is_python3:
builtins.input = raw_input_wrapper
else:
builtins.input = python2_input_wrapper
stdin = StreamStdin(input_queue)
sys.stdin = stdin
class UnbufferedWrite(object):
def __init__(self, stream):
self.stream = stream
def write(self, data):
self.stream.write(data)
self.stream.flush()
def writelines(self, datas):
self.stream.writelines(datas)
self.stream.flush()
def __getattr__(self, attr):
return getattr(self.stream, attr)
def get_stdout_wrapper():
return UnbufferedWrite(open("user_code/user.out", "w"))