<commit_before># coding=utf-8
"""
This module contains sampling related routines
"""

import math
import logging

import numpy as np
import pandas as pd

import sys
from datetime import datetime
import time

import magellan.catalog.catalog_manager as cm

logger = logging.getLogger(__name__)

def _get_stop_words():
    stop_words = [
                "a", "about", "above",
                "across", "after", "afterwards", "again", "against", "all", "almost",
                "alone", "along", "already", "also","although","always","am","among",
                "amongst", "amoungst", "amount", "an", "and", "another", "any", "anyhow",
                "anyone", "anything", "anyway", "anywhere", "are", "around", "as", "at",
                "back","be","became", "because", "become", "becomes", "becoming", "been",
                "before", "beforehand", "behind", "being", "below", "beside", "besides",
                "between", "beyond", "bill", "both", "bottom","but", "by", "call", "can",
                "cannot", "cant", "co", "con", "could", "couldnt", "cry", "de", "describe",
                "detail", "do", "done", "down", "due", "during", "each", "eg", "eight",
                "either", "eleven", "else", "elsewhere", "empty", "enough", "etc", "even",
                "ever", "every", "everyone", "everything", "everywhere", "except", "few",
                "fifteen", "fify", "fill", "find", "fire", "first", "five", "for", "former",
                "formerly", "forty", "found", "four", "from", "front", "full", "further",
                "get", "give", "go", "had", "has", "hasnt", "have", "he", "hence", "her",
                "here", "hereafter", "hereby", "herein", "hereupon", "hers", "herself",
                "him", "himself", "his", "how", "however", "hundred", "ie", "if", "in",
                "inc", "indeed", "interest", "into", "is", "it", "its", "itself", "keep",
                "last", "latter", "latterly", "least", "less", "ltd", "made", "many", "may",
                "me", "meanwhile", "might", "mill", "mine", "more", "moreover", "most",
                "mostly", "move", "much", "must", "my", "myself", "name", "namely",
                "neither", "never", "nevertheless", "next", "nine", "no", "nobody",
                "none", "noone", "nor", "not", "nothing", "now", "nowhere", "of", "off",
                "often", "on", "once", "one", "only", "onto", "or", "other", "others",
                "otherwise", "our", "ours", "ourselves", "out", "over", "own","part",
                "per", "perhaps", "please", "put", "rather", "re", "same", "see", "seem",
                "seemed", "seeming", "seems", "serious", "several", "she", "should",
                "show", "side", "since", "sincere", "six", "sixty", "so", "some",
                "somehow", "someone", "something", "sometime", "sometimes", "somewhere",
                "still", "such", "system", "take", "ten", "than", "that", "the", "their",
                "them", "themselves", "then", "thence", "there", "thereafter", "thereby",
                "therefore", "therein", "thereupon", "these", "they", "thickv", "thin",
                "third", "this", "those", "though", "three", "through", "throughout",
                "thru", "thus", "to", "together", "too", "top", "toward", "towards",
                "twelve", "twenty", "two", "un", "under", "until", "up", "upon", "us",
                "very", "via", "was", "we", "well", "were", "what", "whatever", "when",
                "whence", "whenever", "where", "whereafter", "whereas", "whereby",
                "wherein", "whereupon", "wherever", "whether", "which", "while",
                "whither", "who", "whoever", "whole", "whom", "whose", "why", "will",
                "with", "within", "without", "would", "yet", "you", "your", "yours",
                "yourself", "yourselves"
    ]
    return stop_words

# get string columns
def _get_str_cols(table):
    cols = list(table.columns[table.dtypes==object])
    return cols

# get string column list
def _get_str_cols_list(table):
    cols = list(table.columns[table.dtypes==object])
    col_list = []
    for x in cols:
        col_list.append(table.columns.get_loc(x))
    return col_list

# create inverted index from token to position
def _inv_index(table):
    stop_words = set(_get_stop_words())
    str_cols_ix = _get_str_cols_list(table)
    n = len(table)
    key_pos = dict(zip(range(n), range(n)))
    inv_index = dict()
    pos = 0
    for row in table.itertuples():
        s = ''
        for ix in str_cols_ix:
            s += str(row[ix+1]).lower() + ' '
        s = s.rstrip()
        # tokenize them
        s = set(s.split())
        s = s.difference(stop_words)
        for token in s:
            lst = inv_index.get(token, None)
            if lst is None:
                inv_index[token] = [pos]
            else:
                lst.append(pos)
                inv_index[token] = lst
        pos += 1
    return inv_index

def _probe_index(b_table, y, s_tbl_sz, s_inv_index):
    y_pos = math.floor(y/2)
    h_table = set()
    stop_words = set(_get_stop_words())
    str_cols_ix = _get_str_cols_list(b_table)
    for row in b_table.itertuples():
        s = ''
        for ix in str_cols_ix:
            s += str(row[ix+1]).lower() + ' '
        s = s.rstrip()
        s = set(s.split())
        s = s.difference(stop_words)
        m = set()
        for token in s:
            ids = s_inv_index.get(token, None)
            if ids is not None:
                m.update(ids)
        # pick y/2 elements from m
        k = min(y_pos, len(m))
        m = list(m)
        smpl_pos = np.random.choice(m, k, replace=False)
        s_pos_set = set()
        s_pos_set.update(smpl_pos)
        s_tbl_ids = set(range(s_tbl_sz))
        rem_locs = list(s_tbl_ids.difference(s_pos_set))
        if y - k > 0:
            s_neg_set = np.random.choice(rem_locs, y - k, replace=False)
            h_table.update(s_pos_set, s_neg_set)
    return h_table

# down sample two tables : based on sanjib's index based solution
def down_sample(s_table, b_table, size, y):
    if len(b_table) < size:
        # print('Warning!! size of table B is less than b_size parameter - using entire table B')
        size = len(b_table)

    s_inv_index = _inv_index(s_table)
    b_sample_size = min(math.floor(size/y), len(b_table))
    b_tbl_indices = list(np.random.choice(len(b_table), b_sample_size, replace=False))

    s_tbl_indices = _probe_index(b_table.ix[b_tbl_indices], y,
                                 len(s_table), s_inv_index)
    s_tbl_indices = list(s_tbl_indices)
    l_sampled = s_table.iloc[list(s_tbl_indices)]
    r_sampled = b_table.iloc[list(b_tbl_indices)]

    # update catalog
    cm.copy_properties(s_table, l_sampled)
    cm.copy_properties(b_table, r_sampled)

    return l_sampled, r_sampled

# sample one table using random sampling
def sample_table(table, size, replace=False):
    """
    This is a sample table function
    Args:
        table: data frame
        size: size
        replace: boolean

    Returns:
        sampled table

    """
    if not isinstance(table, pd.DataFrame):
        logger.error('Input table is not of type pandas dataframe')
        raise AssertionError('Input table is not of type pandas dataframe')

    if len(table) == 0:
        logger.error('Size of the input table is 0')
        raise AssertionError('Size of the input table is 0')

    if len(table) < size:
        logger.error('Sample size is larger than the input table size')
        raise AssertionError('Sample size is larger than the input table size')

    s_indices = np.random.choice(len(table), size, replace=replace)
    # sort the indices - just to have an order
    s_indices = sorted(s_indices)
    sampled_table = table.iloc[list(s_indices)]

    cm.init_properties(sampled_table)
    cm.copy_properties(table, sampled_table)
    return sampled_table
<commit_msg>adding comments<commit_after># coding=utf-8
"""
This module contains sampling related routines
"""

import math
import logging

import numpy as np
import pandas as pd

import sys
from datetime import datetime
import time

import magellan.catalog.catalog_manager as cm

logger = logging.getLogger(__name__)

def _get_stop_words():
    stop_words = [
                "a", "about", "above",
                "across", "after", "afterwards", "again", "against", "all", "almost",
                "alone", "along", "already", "also","although","always","am","among",
                "amongst", "amoungst", "amount", "an", "and", "another", "any", "anyhow",
                "anyone", "anything", "anyway", "anywhere", "are", "around", "as", "at",
                "back","be","became", "because", "become", "becomes", "becoming", "been",
                "before", "beforehand", "behind", "being", "below", "beside", "besides",
                "between", "beyond", "bill", "both", "bottom","but", "by", "call", "can",
                "cannot", "cant", "co", "con", "could", "couldnt", "cry", "de", "describe",
                "detail", "do", "done", "down", "due", "during", "each", "eg", "eight",
                "either", "eleven", "else", "elsewhere", "empty", "enough", "etc", "even",
                "ever", "every", "everyone", "everything", "everywhere", "except", "few",
                "fifteen", "fify", "fill", "find", "fire", "first", "five", "for", "former",
                "formerly", "forty", "found", "four", "from", "front", "full", "further",
                "get", "give", "go", "had", "has", "hasnt", "have", "he", "hence", "her",
                "here", "hereafter", "hereby", "herein", "hereupon", "hers", "herself",
                "him", "himself", "his", "how", "however", "hundred", "ie", "if", "in",
                "inc", "indeed", "interest", "into", "is", "it", "its", "itself", "keep",
                "last", "latter", "latterly", "least", "less", "ltd", "made", "many", "may",
                "me", "meanwhile", "might", "mill", "mine", "more", "moreover", "most",
                "mostly", "move", "much", "must", "my", "myself", "name", "namely",
                "neither", "never", "nevertheless", "next", "nine", "no", "nobody",
                "none", "noone", "nor", "not", "nothing", "now", "nowhere", "of", "off",
                "often", "on", "once", "one", "only", "onto", "or", "other", "others",
                "otherwise", "our", "ours", "ourselves", "out", "over", "own","part",
                "per", "perhaps", "please", "put", "rather", "re", "same", "see", "seem",
                "seemed", "seeming", "seems", "serious", "several", "she", "should",
                "show", "side", "since", "sincere", "six", "sixty", "so", "some",
                "somehow", "someone", "something", "sometime", "sometimes", "somewhere",
                "still", "such", "system", "take", "ten", "than", "that", "the", "their",
                "them", "themselves", "then", "thence", "there", "thereafter", "thereby",
                "therefore", "therein", "thereupon", "these", "they", "thickv", "thin",
                "third", "this", "those", "though", "three", "through", "throughout",
                "thru", "thus", "to", "together", "too", "top", "toward", "towards",
                "twelve", "twenty", "two", "un", "under", "until", "up", "upon", "us",
                "very", "via", "was", "we", "well", "were", "what", "whatever", "when",
                "whence", "whenever", "where", "whereafter", "whereas", "whereby",
                "wherein", "whereupon", "wherever", "whether", "which", "while",
                "whither", "who", "whoever", "whole", "whom", "whose", "why", "will",
                "with", "within", "without", "would", "yet", "you", "your", "yours",
                "yourself", "yourselves"
    ]
    return stop_words

# get string columns
def _get_str_cols(table):
    cols = list(table.columns[table.dtypes==object])
    return cols

# get string column list
def _get_str_cols_list(table):
    cols = list(table.columns[table.dtypes==object])
    col_list = []
    for x in cols:
        col_list.append(table.columns.get_loc(x))
    return col_list

# create inverted index from token to position
def _inv_index(table):
    stop_words = set(_get_stop_words())
    str_cols_ix = _get_str_cols_list(table)
    n = len(table)
    key_pos = dict(zip(range(n), range(n)))
    inv_index = dict()
    pos = 0
    for row in table.itertuples():
        s = ''
        for ix in str_cols_ix:
            s += str(row[ix+1]).lower() + ' '
        s = s.rstrip()
        # tokenize them
        s = set(s.split())
        s = s.difference(stop_words)
        for token in s:
            lst = inv_index.get(token, None)
            if lst is None:
                inv_index[token] = [pos]
            else:
                lst.append(pos)
                inv_index[token] = lst
        pos += 1
    return inv_index

def _probe_index(b_table, y, s_tbl_sz, s_inv_index):
    y_pos = math.floor(y/2)
    h_table = set()
    stop_words = set(_get_stop_words())
    str_cols_ix = _get_str_cols_list(b_table)
    for row in b_table.itertuples():
        s = ''
        for ix in str_cols_ix:
            s += str(row[ix+1]).lower() + ' '
        s = s.rstrip()
        s = set(s.split())
        s = s.difference(stop_words)
        m = set()
        for token in s:
            ids = s_inv_index.get(token, None)
            if ids is not None:
                m.update(ids)
        # pick y/2 elements from m
        k = min(y_pos, len(m))
        m = list(m)
        smpl_pos = np.random.choice(m, k, replace=False)
        s_pos_set = set()
        s_pos_set.update(smpl_pos)
        s_tbl_ids = set(range(s_tbl_sz))
        rem_locs = list(s_tbl_ids.difference(s_pos_set))
        if y - k > 0:
            s_neg_set = np.random.choice(rem_locs, y - k, replace=False)
            h_table.update(s_pos_set, s_neg_set)
    return h_table

# down sample of two tables : based on sanjib's index based solution
def down_sample(s_table, b_table, size, y):
    if len(b_table) < size:
        # print('Warning!! size of table B is less than b_size parameter - using entire table B')
        size = len(b_table)

    s_inv_index = _inv_index(s_table)
    b_sample_size = min(math.floor(size/y), len(b_table))
    b_tbl_indices = list(np.random.choice(len(b_table), b_sample_size, replace=False))

    s_tbl_indices = _probe_index(b_table.ix[b_tbl_indices], y,
                                 len(s_table), s_inv_index)
    s_tbl_indices = list(s_tbl_indices)
    l_sampled = s_table.iloc[list(s_tbl_indices)]
    r_sampled = b_table.iloc[list(b_tbl_indices)]

    # update catalog
    cm.copy_properties(s_table, l_sampled)
    cm.copy_properties(b_table, r_sampled)

    return l_sampled, r_sampled

# sample one table using random sampling
def sample_table(table, size, replace=False):
    """
    This is a sample table function
    Args:
        table: data frame
        size: size
        replace: boolean

    Returns:
        sampled table

    """
    if not isinstance(table, pd.DataFrame):
        logger.error('Input table is not of type pandas dataframe')
        raise AssertionError('Input table is not of type pandas dataframe')

    if len(table) == 0:
        logger.error('Size of the input table is 0')
        raise AssertionError('Size of the input table is 0')

    if len(table) < size:
        logger.error('Sample size is larger than the input table size')
        raise AssertionError('Sample size is larger than the input table size')

    s_indices = np.random.choice(len(table), size, replace=replace)
    # sort the indices - just to have an order
    s_indices = sorted(s_indices)
    sampled_table = table.iloc[list(s_indices)]

    cm.init_properties(sampled_table)
    cm.copy_properties(table, sampled_table)
    return sampled_table
<|endoftext|>


<commit_before>"""
Open compressed files transparently.
"""
from __future__ import print_function, division, absolute_import

import gzip
import sys
import io
import os
import time
from subprocess import Popen, PIPE

__version__ = '0.2.1'

_PY3 = sys.version > '3'

try:
	import bz2
except ImportError:
	bz2 = None

try:
	import lzma
except ImportError:
	lzma = None

if _PY3:
	basestring = str

class Closing(object):
	"""
	Inherit from this class and implement a close() method to offer context
	manager functionality.
	"""
	def __enter__(self):
		return self

	def __exit__(self, *exc_info):
		self.close()

	def __del__(self):
		try:
			self.close()
		except:
			pass

class PipedGzipWriter(Closing):
	"""
	Write gzip-compressed files by running an external gzip or pigz process and
	piping into it. On Python 2, this is faster than using gzip.open(). On
	Python 3, it allows to run the compression in a separate process and can
	therefore also be faster.
	"""

	def __init__(self, path, mode='wt'):
		if mode not in ('w', 'wt', 'wb', 'a', 'at', 'ab'):
			raise ValueError("Mode is '{0}', but it must be 'w', 'wt', 'wb', 'a', 'at' or 'ab'".format(mode))
		self.outfile = open(path, mode)
		self.devnull = open(os.devnull, mode)
		self.closed = False

		# Setting close_fds to True in the Popen arguments is necessary due to
		# <http://bugs.python.org/issue12786>.
		kwargs = dict(stdin=PIPE, stdout=self.outfile, stderr=self.devnull, close_fds=True)
		try:
			self.process = Popen(['pigz'], **kwargs)
			self.program = 'pigz'
		except OSError as e:
			# pigz not found, try regular gzip
			try:
				self.process = Popen(['gzip'], **kwargs)
				self.program = 'gzip'
			except (IOError, OSError) as e:
				self.outfile.close()
				self.devnull.close()
				raise
		except IOError as e:
			self.outfile.close()
			self.devnull.close()
			raise
		if _PY3 and 'b' not in mode:
			self._file = io.TextIOWrapper(self.process.stdin)
		else:
			self._file = self.process.stdin

	def write(self, arg):
		self._file.write(arg)

	def close(self):
		self.closed = True
		self._file.close()
		retcode = self.process.wait()
		self.outfile.close()
		self.devnull.close()
		if retcode != 0:
			raise IOError("Output {0} process terminated with exit code {1}".format(self.program, retcode))

class PipedGzipReader(Closing):
	def __init__(self, path, mode='r'):
		if mode not in ('r', 'rt', 'rb'):
			raise ValueError("Mode is '{0}', but it must be 'r', 'rt' or 'rb'".format(mode))
		self.process = Popen(['gzip', '-cd', path], stdout=PIPE, stderr=PIPE)
		if _PY3 and not 'b' in mode:
			self._file = io.TextIOWrapper(self.process.stdout)
		else:
			self._file = self.process.stdout
		if _PY3:
			self._stderr = io.TextIOWrapper(self.process.stderr)
		else:
			self._stderr = self.process.stderr
		self.closed = False
		# Give gzip a little bit of time to report any errors (such as
		# a non-existing file)
		time.sleep(0.01)
		self._raise_if_error()

	def close(self):
		self.closed = True
		retcode = self.process.poll()
		if retcode is None:
			# still running
			self.process.terminate()
		self._raise_if_error()

	def __iter__(self):
		for line in self._file:
			yield line
		self.process.wait()
		self._raise_if_error()

	def _raise_if_error(self):
		"""
		Raise IOError if process is not running anymore and the
		exit code is nonzero.
		"""
		retcode = self.process.poll()
		if retcode is not None and retcode != 0:
			message = self._stderr.read().strip()
			raise IOError(message)

	def read(self, *args):
		data = self._file.read(*args)
		if len(args) == 0 or args[0] <= 0:
			# wait for process to terminate until we check the exit code
			self.process.wait()
		self._raise_if_error()

if bz2 is not None:
	class ClosingBZ2File(bz2.BZ2File, Closing):
		"""
		A better BZ2File that supports the context manager protocol.
		This is relevant only for Python 2.6.
		"""

def xopen(filename, mode='r'):
	"""
	Replacement for the "open" function that can also open files that have
	been compressed with gzip, bzip2 or xz. If the filename is '-', standard
	output (mode 'w') or input (mode 'r') is returned. If the filename ends
	with .gz, the file is opened with a pipe to the gzip program. If that
	does not work, then gzip.open() is used (the gzip module is slower than
	the pipe to the gzip program). If the filename ends with .bz2, it's
	opened as a bz2.BZ2File. Otherwise, the regular open() is used.

	mode can be: 'rt', 'rb', 'at', 'ab', 'wt', or 'wb'
	Instead of 'rt', 'wt' and 'at', 'r', 'w' and 'a' can be used as
	abbreviations.

	In Python 2, the 't' and 'b' characters are ignored.

	Append mode ('a', 'at', 'ab') is unavailable with BZ2 compression and
	will raise an error.
	"""
	if mode in ('r', 'w', 'a'):
		mode += 't'
	if mode not in ('rt', 'rb', 'wt', 'wb', 'at', 'ab'):
		raise ValueError("mode '{0}' not supported".format(mode))
	if not _PY3:
		mode = mode[0]
	if not isinstance(filename, basestring):
		raise ValueError("the filename must be a string")

	# standard input and standard output handling
	if filename == '-':
		return dict(
			r=sys.stdin,
			rt=sys.stdin,
			rb=sys.stdin.buffer,
			w=sys.stdout,
			wt=sys.stdout,
			wb=sys.stdout.buffer)[mode]

	if filename.endswith('.bz2'):
		if bz2 is None:
			raise ImportError("Cannot open bz2 files: The bz2 module is not available")
		if _PY3:
			return bz2.open(filename, mode)
		else:
			if mode[0] == 'a':
				raise ValueError("mode '{0}' not supported with BZ2 compression".format(mode))
			if sys.version_info[:2] <= (2, 6):
				return ClosingBZ2File(filename, mode)
			else:
				return bz2.BZ2File(filename, mode)
	elif filename.endswith('.xz'):
		if lzma is None:
			raise ImportError("Cannot open xz files: The lzma module is not available (use Python 3.3 or newer)")
		return lzma.open(filename, mode)
	elif filename.endswith('.gz'):
		if _PY3:
			return gzip.open(filename, mode, compresslevel=6)
		if sys.version_info[:2] == (2, 7):
			buffered_reader = io.BufferedReader
			buffered_writer = io.BufferedWriter
		else:
			buffered_reader = lambda x: x
			buffered_writer = lambda x: x
		if 'r' in mode:
			try:
				return PipedGzipReader(filename, mode)
			except OSError:
				# gzip not installed
				return buffered_reader(gzip.open(filename, mode))
		else:
			try:
				return PipedGzipWriter(filename, mode)
			except OSError:
				return buffered_writer(gzip.open(filename, mode))
	else:
		return open(filename, mode)
<commit_msg>Add a gzip compression level parameter<commit_after>"""
Open compressed files transparently.
"""
from __future__ import print_function, division, absolute_import

import gzip
import sys
import io
import os
import time
from subprocess import Popen, PIPE

__version__ = '0.2.1'

_PY3 = sys.version > '3'

try:
	import bz2
except ImportError:
	bz2 = None

try:
	import lzma
except ImportError:
	lzma = None

if _PY3:
	basestring = str

class Closing(object):
	"""
	Inherit from this class and implement a close() method to offer context
	manager functionality.
	"""
	def __enter__(self):
		return self

	def __exit__(self, *exc_info):
		self.close()

	def __del__(self):
		try:
			self.close()
		except:
			pass

class PipedGzipWriter(Closing):
	"""
	Write gzip-compressed files by running an external gzip or pigz process and
	piping into it. On Python 2, this is faster than using gzip.open(). On
	Python 3, it allows to run the compression in a separate process and can
	therefore also be faster.
	"""

	def __init__(self, path, mode='wt'):
		if mode not in ('w', 'wt', 'wb', 'a', 'at', 'ab'):
			raise ValueError("Mode is '{0}', but it must be 'w', 'wt', 'wb', 'a', 'at' or 'ab'".format(mode))
		self.outfile = open(path, mode)
		self.devnull = open(os.devnull, mode)
		self.closed = False

		# Setting close_fds to True in the Popen arguments is necessary due to
		# <http://bugs.python.org/issue12786>.
		kwargs = dict(stdin=PIPE, stdout=self.outfile, stderr=self.devnull, close_fds=True)
		try:
			self.process = Popen(['pigz'], **kwargs)
			self.program = 'pigz'
		except OSError as e:
			# pigz not found, try regular gzip
			try:
				self.process = Popen(['gzip'], **kwargs)
				self.program = 'gzip'
			except (IOError, OSError) as e:
				self.outfile.close()
				self.devnull.close()
				raise
		except IOError as e:
			self.outfile.close()
			self.devnull.close()
			raise
		if _PY3 and 'b' not in mode:
			self._file = io.TextIOWrapper(self.process.stdin)
		else:
			self._file = self.process.stdin

	def write(self, arg):
		self._file.write(arg)

	def close(self):
		self.closed = True
		self._file.close()
		retcode = self.process.wait()
		self.outfile.close()
		self.devnull.close()
		if retcode != 0:
			raise IOError("Output {0} process terminated with exit code {1}".format(self.program, retcode))

class PipedGzipReader(Closing):
	def __init__(self, path, mode='r'):
		if mode not in ('r', 'rt', 'rb'):
			raise ValueError("Mode is '{0}', but it must be 'r', 'rt' or 'rb'".format(mode))
		self.process = Popen(['gzip', '-cd', path], stdout=PIPE, stderr=PIPE)
		if _PY3 and not 'b' in mode:
			self._file = io.TextIOWrapper(self.process.stdout)
		else:
			self._file = self.process.stdout
		if _PY3:
			self._stderr = io.TextIOWrapper(self.process.stderr)
		else:
			self._stderr = self.process.stderr
		self.closed = False
		# Give gzip a little bit of time to report any errors (such as
		# a non-existing file)
		time.sleep(0.01)
		self._raise_if_error()

	def close(self):
		self.closed = True
		retcode = self.process.poll()
		if retcode is None:
			# still running
			self.process.terminate()
		self._raise_if_error()

	def __iter__(self):
		for line in self._file:
			yield line
		self.process.wait()
		self._raise_if_error()

	def _raise_if_error(self):
		"""
		Raise IOError if process is not running anymore and the
		exit code is nonzero.
		"""
		retcode = self.process.poll()
		if retcode is not None and retcode != 0:
			message = self._stderr.read().strip()
			raise IOError(message)

	def read(self, *args):
		data = self._file.read(*args)
		if len(args) == 0 or args[0] <= 0:
			# wait for process to terminate until we check the exit code
			self.process.wait()
		self._raise_if_error()

if bz2 is not None:
	class ClosingBZ2File(bz2.BZ2File, Closing):
		"""
		A better BZ2File that supports the context manager protocol.
		This is relevant only for Python 2.6.
		"""

def xopen(filename, mode='r', compresslevel=6):
	"""
	Replacement for the "open" function that can also open files that have
	been compressed with gzip, bzip2 or xz. If the filename is '-', standard
	output (mode 'w') or input (mode 'r') is returned. If the filename ends
	with .gz, the file is opened with a pipe to the gzip program. If that
	does not work, then gzip.open() is used (the gzip module is slower than
	the pipe to the gzip program). If the filename ends with .bz2, it's
	opened as a bz2.BZ2File. Otherwise, the regular open() is used.

	mode can be: 'rt', 'rb', 'at', 'ab', 'wt', or 'wb'
	Instead of 'rt', 'wt' and 'at', 'r', 'w' and 'a' can be used as
	abbreviations.

	In Python 2, the 't' and 'b' characters are ignored.

	Append mode ('a', 'at', 'ab') is unavailable with BZ2 compression and
	will raise an error.

	compresslevel is the gzip compression level. It is not used for bz2 and xz.
	"""
	if mode in ('r', 'w', 'a'):
		mode += 't'
	if mode not in ('rt', 'rb', 'wt', 'wb', 'at', 'ab'):
		raise ValueError("mode '{0}' not supported".format(mode))
	if not _PY3:
		mode = mode[0]
	if not isinstance(filename, basestring):
		raise ValueError("the filename must be a string")

	# standard input and standard output handling
	if filename == '-':
		return dict(
			r=sys.stdin,
			rt=sys.stdin,
			rb=sys.stdin.buffer,
			w=sys.stdout,
			wt=sys.stdout,
			wb=sys.stdout.buffer)[mode]

	if filename.endswith('.bz2'):
		if bz2 is None:
			raise ImportError("Cannot open bz2 files: The bz2 module is not available")
		if _PY3:
			return bz2.open(filename, mode)
		else:
			if mode[0] == 'a':
				raise ValueError("mode '{0}' not supported with BZ2 compression".format(mode))
			if sys.version_info[:2] <= (2, 6):
				return ClosingBZ2File(filename, mode)
			else:
				return bz2.BZ2File(filename, mode)
	elif filename.endswith('.xz'):
		if lzma is None:
			raise ImportError("Cannot open xz files: The lzma module is not available (use Python 3.3 or newer)")
		return lzma.open(filename, mode)
	elif filename.endswith('.gz'):
		if _PY3:
			return gzip.open(filename, mode, compresslevel=compresslevel)
		if sys.version_info[:2] == (2, 7):
			buffered_reader = io.BufferedReader
			buffered_writer = io.BufferedWriter
		else:
			buffered_reader = lambda x: x
			buffered_writer = lambda x: x
		if 'r' in mode:
			try:
				return PipedGzipReader(filename, mode)
			except OSError:
				# gzip not installed
				return buffered_reader(gzip.open(filename, mode))
		else:
			try:
				return PipedGzipWriter(filename, mode)
			except OSError:
				return buffered_writer(gzip.open(filename, mode, compresslevel=compresslevel))
	else:
		return open(filename, mode)
<|endoftext|>


<commit_before># -*- coding: utf-8 -*-
"""Test for short_panel and panel sandwich

Created on Fri May 18 13:05:47 2012

Author: Josef Perktold

moved example from main of random_panel
"""

import numpy as np
from numpy.testing import assert_almost_equal
import numpy.testing as npt
import statsmodels.tools.eval_measures as em
from statsmodels.stats.moment_helpers import cov2corr, se_cov
from statsmodels.regression.linear_model import OLS

from statsmodels.sandbox.panel.panel_short import ShortPanelGLS, ShortPanelGLS2
from statsmodels.sandbox.panel.random_panel import PanelSample
import statsmodels.sandbox.panel.correlation_structures as cs
import statsmodels.stats.sandwich_covariance as sw

def assert_maxabs(actual, expected, value):
    npt.assert_array_less(em.maxabs(actual, expected, None), value)

def test_short_panel():
    #this checks that some basic statistical properties are satisfied by the
    #results, not verified results against other packages
    #Note: the ranking of robust bse is different if within=True
    #I added within keyword to PanelSample to be able to use old example
    #if within is False, then there is no within group variation in exog.
    nobs = 100
    nobs_i = 5
    n_groups = nobs // nobs_i
    k_vars = 3

    dgp = PanelSample(nobs, k_vars, n_groups, corr_structure=cs.corr_arma,
                      corr_args=([1], [1., -0.9],), seed=377769, within=False)
    #print 'seed', dgp.seed
    y = dgp.generate_panel()
    noise = y - dgp.y_true

    #test dgp

    dgp_cov_e = np.array(
              [[ 1.    ,  0.9   ,  0.81  ,  0.729 ,  0.6561],
               [ 0.9   ,  1.    ,  0.9   ,  0.81  ,  0.729 ],
               [ 0.81  ,  0.9   ,  1.    ,  0.9   ,  0.81  ],
               [ 0.729 ,  0.81  ,  0.9   ,  1.    ,  0.9   ],
               [ 0.6561,  0.729 ,  0.81  ,  0.9   ,  1.    ]])

    npt.assert_almost_equal(dgp.cov, dgp_cov_e, 13)

    cov_noise = np.cov(noise.reshape(-1,n_groups, order='F'))
    corr_noise = cov2corr(cov_noise)
    npt.assert_almost_equal(corr_noise, dgp.cov, 1)

    #estimate panel model
    mod2 = ShortPanelGLS(y, dgp.exog, dgp.groups)
    res2 = mod2.fit_iterative(2)

    #whitened residual should be uncorrelated
    corr_wresid = np.corrcoef(res2.wresid.reshape(-1,n_groups, order='F'))
    assert_maxabs(corr_wresid, np.eye(5), 0.1)

    #residual should have same correlation as dgp
    corr_resid = np.corrcoef(res2.resid.reshape(-1,n_groups, order='F'))
    assert_maxabs(corr_resid, dgp.cov, 0.1)

    assert_almost_equal(res2.resid.std(),1, decimal=0)

    y_pred = np.dot(mod2.exog, res2.params)
    assert_almost_equal(res2.fittedvalues, y_pred, 13)

    #compare with OLS

    res2_ols = mod2._fit_ols()
    npt.assert_(mod2.res_pooled is res2_ols)

    res2_ols = mod2.res_pooled  #TODO: BUG: requires call to _fit_ols

    #fitting once is the same as OLS
    #note: I need to create new instance, otherwise it continuous fitting
    mod1 = ShortPanelGLS(y, dgp.exog, dgp.groups)
    res1 = mod1.fit_iterative(1)

    assert_almost_equal(res1.params, res2_ols.params, decimal=13)
    assert_almost_equal(res1.bse, res2_ols.bse, decimal=13)

    res_ols = OLS(y, dgp.exog).fit()
    assert_almost_equal(res1.params, res_ols.params, decimal=13)
    assert_almost_equal(res1.bse, res_ols.bse, decimal=13)

    #compare with old version
    mod_old = ShortPanelGLS2(y, dgp.exog, dgp.groups)
    res_old = mod_old.fit()

    assert_almost_equal(res2.params, res_old.params, decimal=13)
    assert_almost_equal(res2.bse, res_old.bse, decimal=13)

    mod5 = ShortPanelGLS(y, dgp.exog, dgp.groups)
    res5 = mod5.fit_iterative(5)

    #make sure it's different
    #npt.assert_array_less(0.009, em.maxabs(res5.bse, res2.bse))

    cov_clu = sw.cov_cluster(mod2.res_pooled, dgp.groups.astype(int))
    clubse = se_cov(cov_clu)
    pnwbse = se_cov(sw.cov_nw_panel(mod2.res_pooled, 4, mod2.group.groupidx))
    bser = np.vstack((res2.bse, res5.bse, clubse, pnwbse))
    bser_mean = np.mean(bser, axis=0)

    #cov_cluster close to robust and PanelGLS
    #is up to 24% larger than mean of bser
    #npt.assert_array_less(0, clubse / bser_mean - 1)
    npt.assert_array_less(clubse / bser_mean - 1, 0.25)
    #cov_nw_panel close to robust and PanelGLS
    npt.assert_array_less(pnwbse / bser_mean - 1, 0.1)
    #OLS underestimates bse, robust at least 60% larger
    npt.assert_array_less(0.6, bser_mean / res_ols.bse  - 1)

    #cov_hac_panel with uniform_kernel is the same as cov_cluster for balanced
    #panel with full length kernel
    #I fixe default correction to be equal
    cov_uni = sw.cov_nw_panel(mod2.res_pooled, 4, mod2.group.groupidx,
                              weights_func=sw.weights_uniform)
    assert_almost_equal(cov_uni, cov_clu, decimal=13)

    #without correction
    cov_clu2 = sw.cov_cluster(mod2.res_pooled, dgp.groups.astype(int),
                              use_correction=False)
    cov_uni2 = sw.cov_nw_panel(mod2.res_pooled, 4, mod2.group.groupidx,
                              weights_func=sw.weights_uniform,
                              use_correction=False)
    assert_almost_equal(cov_uni2, cov_clu2, decimal=13)

    cov_white = sw.cov_white_simple(mod2.res_pooled)
    cov_pnw0 = sw.cov_nw_panel(mod2.res_pooled, 0, mod2.group.groupidx,
                              use_correction='hac')
    assert_almost_equal(cov_pnw0, cov_white, decimal=13)
<commit_msg>TST: correct test for change in default<commit_after># -*- coding: utf-8 -*-
"""Test for short_panel and panel sandwich

Created on Fri May 18 13:05:47 2012

Author: Josef Perktold

moved example from main of random_panel
"""

import numpy as np
from numpy.testing import assert_almost_equal
import numpy.testing as npt
import statsmodels.tools.eval_measures as em
from statsmodels.stats.moment_helpers import cov2corr, se_cov
from statsmodels.regression.linear_model import OLS

from statsmodels.sandbox.panel.panel_short import ShortPanelGLS, ShortPanelGLS2
from statsmodels.sandbox.panel.random_panel import PanelSample
import statsmodels.sandbox.panel.correlation_structures as cs
import statsmodels.stats.sandwich_covariance as sw

def assert_maxabs(actual, expected, value):
    npt.assert_array_less(em.maxabs(actual, expected, None), value)

def test_short_panel():
    #this checks that some basic statistical properties are satisfied by the
    #results, not verified results against other packages
    #Note: the ranking of robust bse is different if within=True
    #I added within keyword to PanelSample to be able to use old example
    #if within is False, then there is no within group variation in exog.
    nobs = 100
    nobs_i = 5
    n_groups = nobs // nobs_i
    k_vars = 3

    dgp = PanelSample(nobs, k_vars, n_groups, corr_structure=cs.corr_arma,
                      corr_args=([1], [1., -0.9],), seed=377769, within=False)
    #print 'seed', dgp.seed
    y = dgp.generate_panel()
    noise = y - dgp.y_true

    #test dgp

    dgp_cov_e = np.array(
              [[ 1.    ,  0.9   ,  0.81  ,  0.729 ,  0.6561],
               [ 0.9   ,  1.    ,  0.9   ,  0.81  ,  0.729 ],
               [ 0.81  ,  0.9   ,  1.    ,  0.9   ,  0.81  ],
               [ 0.729 ,  0.81  ,  0.9   ,  1.    ,  0.9   ],
               [ 0.6561,  0.729 ,  0.81  ,  0.9   ,  1.    ]])

    npt.assert_almost_equal(dgp.cov, dgp_cov_e, 13)

    cov_noise = np.cov(noise.reshape(-1,n_groups, order='F'))
    corr_noise = cov2corr(cov_noise)
    npt.assert_almost_equal(corr_noise, dgp.cov, 1)

    #estimate panel model
    mod2 = ShortPanelGLS(y, dgp.exog, dgp.groups)
    res2 = mod2.fit_iterative(2)

    #whitened residual should be uncorrelated
    corr_wresid = np.corrcoef(res2.wresid.reshape(-1,n_groups, order='F'))
    assert_maxabs(corr_wresid, np.eye(5), 0.1)

    #residual should have same correlation as dgp
    corr_resid = np.corrcoef(res2.resid.reshape(-1,n_groups, order='F'))
    assert_maxabs(corr_resid, dgp.cov, 0.1)

    assert_almost_equal(res2.resid.std(),1, decimal=0)

    y_pred = np.dot(mod2.exog, res2.params)
    assert_almost_equal(res2.fittedvalues, y_pred, 13)

    #compare with OLS

    res2_ols = mod2._fit_ols()
    npt.assert_(mod2.res_pooled is res2_ols)

    res2_ols = mod2.res_pooled  #TODO: BUG: requires call to _fit_ols

    #fitting once is the same as OLS
    #note: I need to create new instance, otherwise it continuous fitting
    mod1 = ShortPanelGLS(y, dgp.exog, dgp.groups)
    res1 = mod1.fit_iterative(1)

    assert_almost_equal(res1.params, res2_ols.params, decimal=13)
    assert_almost_equal(res1.bse, res2_ols.bse, decimal=13)

    res_ols = OLS(y, dgp.exog).fit()
    assert_almost_equal(res1.params, res_ols.params, decimal=13)
    assert_almost_equal(res1.bse, res_ols.bse, decimal=13)

    #compare with old version
    mod_old = ShortPanelGLS2(y, dgp.exog, dgp.groups)
    res_old = mod_old.fit()

    assert_almost_equal(res2.params, res_old.params, decimal=13)
    assert_almost_equal(res2.bse, res_old.bse, decimal=13)

    mod5 = ShortPanelGLS(y, dgp.exog, dgp.groups)
    res5 = mod5.fit_iterative(5)

    #make sure it's different
    #npt.assert_array_less(0.009, em.maxabs(res5.bse, res2.bse))

    cov_clu = sw.cov_cluster(mod2.res_pooled, dgp.groups.astype(int))
    clubse = se_cov(cov_clu)
    pnwbse = se_cov(sw.cov_nw_panel(mod2.res_pooled, 4, mod2.group.groupidx))
    bser = np.vstack((res2.bse, res5.bse, clubse, pnwbse))
    bser_mean = np.mean(bser, axis=0)

    #cov_cluster close to robust and PanelGLS
    #is up to 24% larger than mean of bser
    #npt.assert_array_less(0, clubse / bser_mean - 1)
    npt.assert_array_less(clubse / bser_mean - 1, 0.25)
    #cov_nw_panel close to robust and PanelGLS
    npt.assert_array_less(pnwbse / bser_mean - 1, 0.1)
    #OLS underestimates bse, robust at least 60% larger
    npt.assert_array_less(0.6, bser_mean / res_ols.bse  - 1)

    #cov_hac_panel with uniform_kernel is the same as cov_cluster for balanced
    #panel with full length kernel
    #I fixe default correction to be equal
    cov_uni = sw.cov_nw_panel(mod2.res_pooled, 4, mod2.group.groupidx,
                              weights_func=sw.weights_uniform,
                              use_correction='c')
    assert_almost_equal(cov_uni, cov_clu, decimal=13)

    #without correction
    cov_clu2 = sw.cov_cluster(mod2.res_pooled, dgp.groups.astype(int),
                              use_correction=False)
    cov_uni2 = sw.cov_nw_panel(mod2.res_pooled, 4, mod2.group.groupidx,
                              weights_func=sw.weights_uniform,
                              use_correction=False)
    assert_almost_equal(cov_uni2, cov_clu2, decimal=13)

    cov_white = sw.cov_white_simple(mod2.res_pooled)
    cov_pnw0 = sw.cov_nw_panel(mod2.res_pooled, 0, mod2.group.groupidx,
                              use_correction='hac')
    assert_almost_equal(cov_pnw0, cov_white, decimal=13)
<|endoftext|>


<commit_before>#! /usr/bin/env python3

# Copyright 2018 John Hanley.
#
# Permission is hereby granted, free of charge, to any person obtaining a
# copy of this software and associated documentation files (the "Software"),
# to deal in the Software without restriction, including without limitation
# the rights to use, copy, modify, merge, publish, distribute, sublicense,
# and/or sell copies of the Software, and to permit persons to whom the
# Software is furnished to do so, subject to the following conditions:
# The above copyright notice and this permission notice shall be included in
# all copies or substantial portions of the Software.
# The software is provided "AS IS", without warranty of any kind, express or
# implied, including but not limited to the warranties of merchantability,
# fitness for a particular purpose and noninfringement. In no event shall
# the authors or copyright holders be liable for any claim, damages or
# other liability, whether in an action of contract, tort or otherwise,
# arising from, out of or in connection with the software or the use or
# other dealings in the software.

import itertools
import random
import unittest

from bs4 import BeautifulSoup
import requests

def merge_streams(a, b):
    """Given pre-sorted iterables A and B, yields their sorted union.
    """
    a_item, valid = _get_next(a)
    if not valid:  # Input A was empty.
        yield from b
        return

    b_item, valid = _get_next(a)
    if not valid:  # Input B was empty.
        yield a_item
        yield from a
        return

    # Loop invariant:
    #   We have a & b items we can compare, and
    #   both streams still potentially have items to consume.

    while True:

        if a_item <= b_item:
            yield a_item
            a_item, valid = _get_next(a)
            if not valid:  # We found the end of A.
                yield b_item
                yield from b
                return

        else:

            yield b_item
            b_item, valid = _get_next(b)
            if not valid:  # We found the end of B.
                yield a_item
                yield from a
                return

def _get_next(it):
    try:
        return next(it), True
    except StopIteration as e:
        return None, False

# From https://docs.python.org/3/library/itertools.html#itertools-recipes
def take(n, iterable):
    "Return first n items of the iterable as a list"
    return list(itertools.islice(iterable, n))

class ItemCountGenerator:
    """Produces an order sequence of (item, count) tuples, as a testing aid.
    """

    def __init__(self, num_valid_items=3, seed=None):
        k = num_valid_items
        assert k <= 20
        self.valid_items = take(k, self._get_nato_letters())
        # self.valid_items = take(k, self._get_valid_items())
        if seed:
            random.seed(seed)

    def generate(self):
        while True:
            i = random.randrange(len(self.valid_items))
            yield self.valid_items[i]

    @staticmethod
    def _vowel_count(s):
        return sum(map(s.count, 'aeiouy'))

    def _get_valid_items(self, words='/usr/share/dict/web2a', size=4):
        next = '`'  # next candidate prefix
        with open(words) as fin:
            for line in fin:
                if line <= next:
                    continue
                word = line.split()[0].lower()
                if (len(word) == size
                        and self._vowel_count(word) == 1):
                    yield word
                    next = chr(1 + ord(word[0]))

    def _get_nato_letters(self):
        url = 'https://en.wiktionary.org/wiki/Appendix:NATO_phonetic_alphabet'
        r = requests.get(url)
        tbl = BeautifulSoup(r.text, 'html5lib').find('table')
        for row in tbl.find_all('tr'):
            tds = list(row.find_all('td'))
            if tds and len(tds[0].text) == 1:
                letter_text = tds[1].text.lower()
                yield letter_text.split()[-1]  # This handles xray.

class ExternalMergesortTest(unittest.TestCase):

    def test_one_empty(self):
        a = iter('')
        b = iter('xy')
        self.assertEqual('xy', ''.join(merge_streams(a, b)))

        a = iter('xy')
        b = iter('')
        self.assertEqual('xy', ''.join(merge_streams(a, b)))

    def test_interleaved(self):
        a = iter('abxyz')
        b = iter('cdef')
        self.assertEqual('abcdefxyz', ''.join(merge_streams(a, b)))

        a = iter('abddexyz')
        b = iter('ceeff')
        self.assertEqual('abcddeeeffxyz', ''.join(merge_streams(a, b)))

if __name__ == '__main__':
    ic_gen = ItemCountGenerator(seed=42)
    k = 3
    a = iter(take(2 * k, ic_gen.generate()))
    b = iter(take(k, ic_gen.generate()))

    lst = list(merge_streams(a, b))

    unittest.main()
<commit_msg>Accesses the streams via an array index.<commit_after>#! /usr/bin/env python3

# Copyright 2018 John Hanley.
#
# Permission is hereby granted, free of charge, to any person obtaining a
# copy of this software and associated documentation files (the "Software"),
# to deal in the Software without restriction, including without limitation
# the rights to use, copy, modify, merge, publish, distribute, sublicense,
# and/or sell copies of the Software, and to permit persons to whom the
# Software is furnished to do so, subject to the following conditions:
# The above copyright notice and this permission notice shall be included in
# all copies or substantial portions of the Software.
# The software is provided "AS IS", without warranty of any kind, express or
# implied, including but not limited to the warranties of merchantability,
# fitness for a particular purpose and noninfringement. In no event shall
# the authors or copyright holders be liable for any claim, damages or
# other liability, whether in an action of contract, tort or otherwise,
# arising from, out of or in connection with the software or the use or
# other dealings in the software.

import enum
import itertools
import random
import unittest

from bs4 import BeautifulSoup
import requests

class Stream(enum.IntEnum):
    """Stream Identifier, used as item[] index."""
    A = 0
    B = 1

def merge_streams(a, b):
    """Given pre-sorted iterables A and B, yields their sorted union.
    """
    item = [None, None]
    item[Stream.A], valid = _get_next(a)
    if not valid:  # Input A was empty.
        yield from b
        return

    item[Stream.B], valid = _get_next(a)
    if not valid:  # Input B was empty.
        yield item[Stream.A]
        yield from a
        return

    stream = (a, b)

    # Choose the "winning" (or lower valued) stream to consume from.
    a_item, b_item = item
    cur_stream = Stream.A if a_item <= b_item else Stream.B

    # Loop invariant:
    #   We have a & b items we can compare, and
    #   both streams still potentially have items to consume.

    while True:

        yield item[cur_stream]
        item[cur_stream], valid = _get_next(stream[cur_stream])

        if valid:
            a_item, b_item = item
            cur_stream = Stream.A if a_item <= b_item else Stream.B
        else:  # We found the end of this stream.
            alt = 1 - cur_stream  # Let's drain the alternate stream.
            yield item[alt]
            yield from stream[alt]
            return

def _get_next(it):
    try:
        return next(it), True
    except StopIteration as e:
        return None, False

# From https://docs.python.org/3/library/itertools.html#itertools-recipes
def take(n, iterable):
    "Return first n items of the iterable as a list"
    return list(itertools.islice(iterable, n))

class ItemCountGenerator:
    """Produces an order sequence of (item, count) tuples, as a testing aid.
    """

    def __init__(self, num_valid_items=3, seed=None):
        k = num_valid_items
        assert k <= 20
        self.valid_items = take(k, self._get_nato_letters())
        # self.valid_items = take(k, self._get_valid_items())
        if seed:
            random.seed(seed)

    def generate(self):
        while True:
            i = random.randrange(len(self.valid_items))
            yield self.valid_items[i]

    @staticmethod
    def _vowel_count(s):
        return sum(map(s.count, 'aeiouy'))

    def _get_valid_items(self, words='/usr/share/dict/web2a', size=4):
        next = '`'  # next candidate prefix
        with open(words) as fin:
            for line in fin:
                if line <= next:
                    continue
                word = line.split()[0].lower()
                if (len(word) == size
                        and self._vowel_count(word) == 1):
                    yield word
                    next = chr(1 + ord(word[0]))

    def _get_nato_letters(self):
        url = 'https://en.wiktionary.org/wiki/Appendix:NATO_phonetic_alphabet'
        r = requests.get(url)
        tbl = BeautifulSoup(r.text, 'html5lib').find('table')
        for row in tbl.find_all('tr'):
            tds = list(row.find_all('td'))
            if tds and len(tds[0].text) == 1:
                letter_text = tds[1].text.lower()
                yield letter_text.split()[-1]  # This handles xray.

class ExternalMergesortTest(unittest.TestCase):

    def test_one_empty(self):
        a = iter('')
        b = iter('xy')
        self.assertEqual('xy', ''.join(merge_streams(a, b)))

        a = iter('xy')
        b = iter('')
        self.assertEqual('xy', ''.join(merge_streams(a, b)))

    def test_interleaved(self):
        a = iter('abxyz')
        b = iter('cdef')
        self.assertEqual('abcdefxyz', ''.join(merge_streams(a, b)))

        a = iter('abddexyz')
        b = iter('ceeff')
        self.assertEqual('abcddeeeffxyz', ''.join(merge_streams(a, b)))

if __name__ == '__main__':
    ic_gen = ItemCountGenerator(seed=42)
    k = 3
    a = iter(take(2 * k, ic_gen.generate()))
    b = iter(take(k, ic_gen.generate()))

    lst = list(merge_streams(a, b))

    unittest.main()
<|endoftext|>


<commit_before>#!/usr/bin/env python

from src.util import date_string_from_duration_in_seconds, date_string_from_epoch_timestamp

def calculate_compute_time(job, node):
    """Given a Job and ComputeNode, determine how long the job will compute.

    Use DISCOVAR benchmark data.
    """
    # TODO
    return job.historical_end_time - job.historical_start_time

class ComputeNode:

    def __init__(self, name, running_watts, idle_watts, cpus):
        self.name = name
        self.running_watts = running_watts
        self.idle_watts = idle_watts
        self.cpus = cpus
        self.compute_times = []
        self.idle_times = []
        self.current_jobs = []
        self.compute_time = 0
        self.idle_time = 0
        self.current_time = None

    def was_running(self, time):
        for job in self.current_jobs:
            if job.start_time <= time and job.end_time >= time:
                return True
        return False

    def cpus_in_use(self, time):
        count = 0
        for job in self.current_jobs:
            if job.start_time <= time and job.end_time >= time:
                count += job.cpus_requested
        return count

    def x_cpus_available_for_y_seconds(self, cpus_requested, seconds, start_time):
        for time in range(start_time, start_time+seconds+1):
            cpus_available = self.cpus - self.cpus_in_use(time)
            if cpus_available < cpus_requested:
                return False
        return True

    def initialize(self, time):
        self.current_time = time

    def update(self, newtime):
        """Returns a list of jobs completed in the period previous to newtime.

        Also increments self.compute_time or self.idle_time.
        """
        completed_jobs = []
        for moment in range(self.current_time, newtime):
            for job in self.current_jobs:
                if job.end_time <= moment:
                    completed_jobs.append(job)
                    # remove this job from self.current_jobs
                    self.current_jobs = [j for j in self.current_jobs if j not in completed_jobs]
            if self.was_running(moment):
                self.compute_time += 1
            else:
                self.idle_time += 1
        self.current_time = newtime
        return completed_jobs

    def calculate_total_compute_time(self):
        self.compute_times = sorted(self.compute_times)
        wall_times = []
        current_start = None
        current_end = None
        for time_tuple in self.compute_times:
            if not current_start:
                current_start = time_tuple[0]
                current_end = time_tuple[1]
                continue

            if time_tuple[0] > current_end:
                wall_times.append( (current_start, current_end) )
                current_start = time_tuple[0]
                current_end = time_tuple[1]
            elif time_tuple[0] == current_end:
                current_end = time_tuple[1]
            elif time_tuple[0] < current_end:
                if time_tuple[1] > current_end:
                    current_end = time_tuple[1]
        wall_times.append( (current_start, current_end) )
        total_compute_time = sum([t[1] - t[0] for t in wall_times])
        return total_compute_time

    def generate_report(self):
        # name\ttotal_compute_time\ttotal_idle_time\t
        #        total_energy_consumption (kWh)
        running_cost = (self.compute_time * self.running_watts / 3600) / 1000
        idle_cost = (self.idle_time * self.idle_watts / 3600) / 1000
        energy_cost = running_cost + idle_cost
        compute_time = date_string_from_duration_in_seconds(self.compute_time)
        idle_time = date_string_from_duration_in_seconds(self.idle_time)
        data = [self.name, compute_time, idle_time, str(energy_cost)]
        return "\t".join(data) + "\n"

    def add_job(self, job):
        self.current_jobs.append(job)

    def find_job_start_time(self, job):
        """Given a job, find the earliest start time available.

        Consider the cpus requested for the job in question as well as
          any running and scheduled jobs.

        Return the time in epoch seconds.
        """
        TIME_LIMIT = 1000 # only check for 1000 sec after job arrival
        time = self.current_time
        job_start_time = None
        print("find_job_start_time here, looking for a time to start job " + str(job))
        print("currently have " + str(len(self.current_jobs)) + " jobs...")
        print("and using " + str(self.cpus_in_use(time)) + " at time " + str(time))
        while not job_start_time:
            if time % 10 == 0:
                print("\tnow checking time " + str(time))
            if time == job.arrival_time + TIME_LIMIT:
                return None
            cpus_available = self.cpus - self.cpus_in_use(time)
            if cpus_available >= job.cpus_requested:
                if self.x_cpus_available_for_y_seconds(job.cpus_requested, job.compute_time, time):
                    print("\tfound it! the job can run at time " + str(time) + "\n\n")
                    return time
            time += 1
<commit_msg>cleaned up compute node update a bit<commit_after>#!/usr/bin/env python

from src.util import date_string_from_duration_in_seconds, date_string_from_epoch_timestamp

def calculate_compute_time(job, node):
    """Given a Job and ComputeNode, determine how long the job will compute.

    Use DISCOVAR benchmark data.
    """
    # TODO
    return job.historical_end_time - job.historical_start_time

class ComputeNode:

    def __init__(self, name, running_watts, idle_watts, cpus):
        self.name = name
        self.running_watts = running_watts
        self.idle_watts = idle_watts
        self.cpus = cpus
        self.compute_times = []
        self.idle_times = []
        self.current_jobs = []
        self.compute_time = 0
        self.idle_time = 0
        self.current_time = None

    def was_running(self, time):
        for job in self.current_jobs:
            if job.start_time <= time and job.end_time >= time:
                return True
        return False

    def cpus_in_use(self, time):
        count = 0
        for job in self.current_jobs:
            if job.start_time <= time and job.end_time >= time:
                count += job.cpus_requested
        return count

    def x_cpus_available_for_y_seconds(self, cpus_requested, seconds, start_time):
        for time in range(start_time, start_time+seconds+1):
            cpus_available = self.cpus - self.cpus_in_use(time)
            if cpus_available < cpus_requested:
                return False
        return True

    def initialize(self, time):
        self.current_time = time

    def update(self, newtime):
        """Returns a list of jobs completed in the period previous to newtime.

        Also increments self.compute_time or self.idle_time.
        """
        completed_jobs = []
        for job in self.current_jobs:
            if job.end_time < newtime:
                completed_jobs.append(job)
                # remove this job from self.current_jobs
                self.current_jobs = [j for j in self.current_jobs if j not in completed_jobs]
        # increment running/idle time
        if self.was_running(newtime-1):
            self.compute_time += 1
        else:
            self.idle_time += 1
        self.current_time = newtime
        return completed_jobs

    def calculate_total_compute_time(self):
        self.compute_times = sorted(self.compute_times)
        wall_times = []
        current_start = None
        current_end = None
        for time_tuple in self.compute_times:
            if not current_start:
                current_start = time_tuple[0]
                current_end = time_tuple[1]
                continue

            if time_tuple[0] > current_end:
                wall_times.append( (current_start, current_end) )
                current_start = time_tuple[0]
                current_end = time_tuple[1]
            elif time_tuple[0] == current_end:
                current_end = time_tuple[1]
            elif time_tuple[0] < current_end:
                if time_tuple[1] > current_end:
                    current_end = time_tuple[1]
        wall_times.append( (current_start, current_end) )
        total_compute_time = sum([t[1] - t[0] for t in wall_times])
        return total_compute_time

    def generate_report(self):
        # name\ttotal_compute_time\ttotal_idle_time\t
        #        total_energy_consumption (kWh)
        running_cost = (self.compute_time * self.running_watts / 3600) / 1000
        idle_cost = (self.idle_time * self.idle_watts / 3600) / 1000
        energy_cost = running_cost + idle_cost
        compute_time = date_string_from_duration_in_seconds(self.compute_time)
        idle_time = date_string_from_duration_in_seconds(self.idle_time)
        data = [self.name, compute_time, idle_time, str(energy_cost)]
        return "\t".join(data) + "\n"

    def add_job(self, job):
        self.current_jobs.append(job)

    def find_job_start_time(self, job):
        """Given a job, find the earliest start time available.

        Consider the cpus requested for the job in question as well as
          any running and scheduled jobs.

        Return the time in epoch seconds.
        """
        TIME_LIMIT = 1000 # only check for 1000 sec after job arrival
        time = self.current_time
        job_start_time = None
        print("find_job_start_time here, looking for a time to start job " + str(job))
        print("currently have " + str(len(self.current_jobs)) + " jobs...")
        print("and using " + str(self.cpus_in_use(time)) + " at time " + str(time))
        while not job_start_time:
            if time % 10 == 0:
                print("\tnow checking time " + str(time))
            if time == job.arrival_time + TIME_LIMIT:
                return None
            cpus_available = self.cpus - self.cpus_in_use(time)
            if cpus_available >= job.cpus_requested:
                if self.x_cpus_available_for_y_seconds(job.cpus_requested, job.compute_time, time):
                    print("\tfound it! the job can run at time " + str(time) + "\n\n")
                    return time
            time += 1
<|endoftext|>


<commit_before>import base64
import falcon
import json
import logging

from engine.runners.python_runner import PythonRunner
import engine.util as util

# Config applies to all other loggers
logging.basicConfig(format='[%(asctime)s] %(name)s:%(levelname)s: %(message)s', level=logging.DEBUG)
logger = logging.getLogger(__name__)

class SubmitResource(object):

    def on_post(self, req, resp):
        """Handle POST requests."""
        try:
            raw_json = req.stream.read()
        except Exception as ex:
            raise falcon.HTTPError(falcon.HTTP_400, 'Error', ex.message)

        try:
            result_json = json.loads(raw_json.decode('utf-8'))
        except ValueError:
            raise falcon.HTTPError(falcon.HTTP_400, 'Invalid JSON', 'Could not decode the request body.')

        # Write user's code to file.
        code = result_json.get('code')
        if code:
            decoded_code = str(base64.b64decode(code), 'utf-8')
            code_filename = util.write_str_to_file(decoded_code)
            logger.debug("User code saved in: %s", code_filename)
        else:
            raise falcon.HTTPError(falcon.HTTP_400, 'Invalid JSON.', 'No code provided.')

        # Create a Problem1 instance which will come generated with several test cases.
        # TODO: Find a better way of deciding which problem to create an instance of.
        problem_id = int(result_json.get('problemID'))
        if problem_id == 1:
            import problems.problem1 as problem
        else:
            raise falcon.HTTPError(falcon.HTTP_400, 'Invalid JSON.', 'Invalid problem ID.')

        # Generate a bunch of test cases.from
        test_cases = problem.generate_test_cases()
        n_cases = len(test_cases)

        n = 1  # test case counter
        passes = [None]*n_cases
        details_html = '
'
        test_case_details = ''

        for tc in test_cases:
            input_str = tc.input_str()
            logger.info("Test case %d/%d (%s).", n, n_cases, tc.test_type.test_name)
            tc_panel_title = 'Test case {:d}/{:d} ({}): '.format(n, n_cases, tc.test_type.test_name)

            logger.debug("Input string:")
            logger.debug("%s", input_str)

            user_answer, process_info = PythonRunner().run(code_filename, input_str)

            logger.debug("Output string:")
            logger.debug("%s", user_answer)

            p_info_str = 'Return code {:d}, utime: {:f} s, stime: {:f} s, maxrss: {:d} kB.
'\
                .format(process_info['return_code'], process_info['utime'], process_info['stime'],
                        process_info['maxrss'])  # resource
            #           process_info['rss'])  # psutil

            passed = problem.verify_user_solution(input_str, user_answer)

            if passed:
                passes[n-1] = True
                logger.info("Test case passed!")
                tc_panel_title += 'passed!'
            else:
                passes[n-1] = False
                logger.info("Test case failed.")
                tc_panel_title += 'failed.'

            test_case_details += '
'\
                                 + '
' + tc_panel_title + '
'\
                                 + '
'\
                                 + p_info_str\
                                 + 'Input:
' + input_str + '

'\
                                 + 'Output:
' + user_answer + '
'\
                                 + '
'
            n = n+1

        n_passes = sum(passes)  # Turns out you can sum booleans (True=1, False=0).

        logger.info("Passed %d/%d test cases.", n_passes, n_cases)
        details_html += 'Passed {:d}/{:d} test cases.
'.format(n_passes, n_cases)
        details_html += test_case_details + '
'

        all_solved = True if n_passes == n_cases else False

        resp_dict = {
            'details': details_html,
            'success': True if all_solved else False
        }

        resp.status = falcon.HTTP_200
        resp.set_header('Access-Control-Allow-Origin', '*')
        resp.body = json.dumps(resp_dict)

        util.delete_file(code_filename)
        logger.debug("User code file deleted: %s", code_filename)

app = falcon.API()
app.add_route('/submit', SubmitResource())
<commit_msg>Engine can now import arbitrary problems.<commit_after>import base64
import falcon
import json
import logging
import importlib

from engine.runners.python_runner import PythonRunner
import engine.util as util

# Config applies to all other loggers
logging.basicConfig(format='[%(asctime)s] %(name)s:%(levelname)s: %(message)s', level=logging.DEBUG)
logger = logging.getLogger(__name__)

class SubmitResource(object):

    def on_post(self, req, resp):
        """Handle POST requests."""
        try:
            raw_json = req.stream.read()
        except Exception as ex:
            raise falcon.HTTPError(falcon.HTTP_400, 'Error', ex.message)

        try:
            result_json = json.loads(raw_json.decode('utf-8'))
        except ValueError:
            raise falcon.HTTPError(falcon.HTTP_400, 'Invalid JSON', 'Could not decode the request body.')

        # Write user's code to file.
        code = result_json.get('code')
        if code:
            decoded_code = str(base64.b64decode(code), 'utf-8')
            code_filename = util.write_str_to_file(decoded_code)
            logger.debug("User code saved in: %s", code_filename)
        else:
            raise falcon.HTTPError(falcon.HTTP_400, 'Invalid JSON.', 'No code provided.')

        # Fetch problem ID and load the correct problem module.
        problem_id = result_json.get('problemID')
        try:
            problem = importlib.import_module('problems.' + problem_id)
            test_case_type_enum = problem.TEST_CASE_TYPE_ENUM
        except Exception:
            logger.error("Could not import module `problems.%s`", problem_id)
            logger.error("Returning HTTP 400 Bad Request due to possibly invalid JSON.")
            raise falcon.HTTPError(falcon.HTTP_400, 'Invalid JSON.', 'Invalid problemID!')

        logger.info("Generating test cases...")
        test_cases = []

        # Count number of test cases we'll be generating in total.
        n_cases = 0
        for test_type in test_case_type_enum:
            n_cases += test_type.multiplicity

        # Generate all the cases and store them in test_cases.
        n = 1
        for test_type in test_case_type_enum:
            for _ in range(test_type.multiplicity):
                logger.debug("Generating test case %d/%d...", n, n_cases)
                test_cases.append(problem.generate_input(test_type))
                n = n+1

        n = 1  # test case counter
        passes = n_cases*[None]
        details_html = '
'
        test_case_details = ''

        for tc in test_cases:
            input_str = tc.input_str()
            logger.info("Test case %d/%d (%s).", n, n_cases, tc.test_type.test_name)
            tc_panel_title = 'Test case {:d}/{:d} ({}): '.format(n, n_cases, tc.test_type.test_name)

            logger.debug("Input string:")
            logger.debug("%s", input_str)

            user_answer, process_info = PythonRunner().run(code_filename, input_str)

            logger.debug("Output string:")
            logger.debug("%s", user_answer)

            p_info_str = 'Return code {:d}, utime: {:f} s, stime: {:f} s, maxrss: {:d} kB.
'\
                .format(process_info['return_code'], process_info['utime'], process_info['stime'],
                        process_info['maxrss'])  # resource
            #           process_info['rss'])  # psutil

            passed = problem.verify_user_solution(input_str, user_answer)

            if passed:
                passes[n-1] = True
                logger.info("Test case passed!")
                tc_panel_title += 'passed!'
            else:
                passes[n-1] = False
                logger.info("Test case failed.")
                tc_panel_title += 'failed.'

            # TODO: Move this html generation to the client side. Send the raw input/output so the client can do it.
            test_case_details += '
'\
                                 + '
' + tc_panel_title + '
'\
                                 + '
'\
                                 + p_info_str\
                                 + 'Input:
' + input_str + '

'\
                                 + 'Output:
' + user_answer + '
'\
                                 + '
'
            n = n+1

        n_passes = sum(passes)  # Turns out you can sum booleans (True=1, False=0).

        logger.info("Passed %d/%d test cases.", n_passes, n_cases)
        details_html += 'Passed {:d}/{:d} test cases.
'.format(n_passes, n_cases)
        details_html += test_case_details + '
'

        all_solved = True if n_passes == n_cases else False

        resp_dict = {
            'details': details_html,
            'success': True if all_solved else False
        }

        resp.status = falcon.HTTP_200
        resp.set_header('Access-Control-Allow-Origin', '*')
        resp.body = json.dumps(resp_dict)

        util.delete_file(code_filename)
        logger.debug("User code file deleted: %s", code_filename)

app = falcon.API()
app.add_route('/submit', SubmitResource())
<|endoftext|>


<commit_before>import os
import json
import base64
import shutil
import urllib
import logging
import datetime
import importlib
import traceback

import falcon

import engine.util as util
from engine.simple_lxd import simple_lxd as lxd
from engine.runners.python_runner import PythonRunner, FilePushError, FilePullError, EngineExecutionError
from engine.runners.javascript_runner import JavascriptRunner
from engine.runners.julia_runner import JuliaRunner

log_file_path = os.path.join(os.path.dirname(os.path.abspath(__file__)), "logging.ini")
logging.config.fileConfig(log_file_path, disable_existing_loggers=False)
logger = logging.getLogger(__name__)

cwd = os.path.dirname(os.path.abspath(__file__))
os.chdir(cwd)

class SubmitResource(object):
    def __init__(self):
        self.pid = os.getpid()
        self.container_image = "lovelace-image"
        self.container_name = "lovelace-{:d}-{:s}".format(self.pid, datetime.datetime.now().strftime("%Y%m%d-%H%M%S"))

        logger.info("Launching Linux container {:s} (pid={:d}) from image {:s}..."
                    .format(self.container_name, self.pid, self.container_image))

        lxd.launch(self.container_image, name=self.container_name, profile="lovelace")

    def on_post(self, req, resp):
        payload = req.media

        code = payload['code']
        language = payload['language']

        if not code:
            resp_dict = {'error': "No code provided!"}
            resp.status = falcon.HTTP_400
            resp.set_header('Access-Control-Allow-Origin', '*')
            resp.body = json.dumps(resp_dict)
            return

        code_filename = write_code_to_file(code, language)

        try:
            # Fetch problem ID and load the correct problem module.
            problem_name = payload['problem'].replace('-', '_')
            problem_module = 'problems.{:s}'.format(problem_name)

            logger.debug("Importing problem_name={:s} problem_module={:s}...".format(problem_name, problem_module))

            problem = importlib.import_module(problem_module)
        except Exception:
            explanation = "Could not import module {:s}. " \
                          "Returning HTTP 400 Bad Request due to possibly invalid JSON.".format(problem_module)
            add_error_to_response(resp, explanation, traceback.format_exc(), falcon.HTTP_400, code_filename)
            return

        function_name = problem.FUNCTION_NAME
        problem_dir = problem_name

        # Copy static resources into engine directory and push them into the Linux container.
        static_resources = []
        for resource_file_name in problem.STATIC_RESOURCES:
            from_path = os.path.join(cwd, "..", "resources", problem_dir, resource_file_name)
            to_path = os.path.join(cwd, resource_file_name)

            logger.debug("Copying static resource from {:s} to {:s}".format(from_path, to_path))

            try:
                shutil.copyfile(from_path, to_path)
            except Exception:
                explanation = "Engine failed to copy a static resource. Returning falcon HTTP 500."
                add_error_to_response(resp, explanation, traceback.format_exc(), falcon.HTTP_500, code_filename)
                return

            static_resources.append(to_path)

            container_path = "/root/{:}".format(resource_file_name)
            logger.debug("Pushing static resource to container {:}{:}".format(self.container_name, container_path))
            lxd.file_push(self.container_name, from_path, container_path)

        logger.info("Generating test cases...")
        test_cases = []

        try:
            for i, test_type in enumerate(problem.TestCaseType):
                for j in range(test_type.multiplicity):
                    logger.debug("Generating test case {:d}: {:s} ({:d}/{:d})..."
                                 .format(len(test_cases)+1, str(test_type), j+1, test_type.multiplicity))
                    test_cases.append(problem.generate_test_case(test_type))
        except Exception:
            explanation = "Engine failed to generate a test case. Returning falcon HTTP 500."
            add_error_to_response(resp, explanation, traceback.format_exc(), falcon.HTTP_500, code_filename)
            return

        # Copy over all the dynamic resources generated by the test cases.
        dynamic_resources = []
        for i, tc in enumerate(test_cases):
            if 'DYNAMIC_RESOURCES' in tc.input:
                for dynamic_resource_filename in tc.input['DYNAMIC_RESOURCES']:
                    resource_path = os.path.join(cwd, "..", "resources", problem_dir, dynamic_resource_filename)
                    destination_path = os.path.join(cwd, dynamic_resource_filename)

                    logger.debug("Copying test case resource from {:s} to {:s}..."
                                 .format(resource_path, destination_path))

                    shutil.copyfile(resource_path, destination_path)
                    dynamic_resources.append(destination_path)

                    container_path = "/root/{:}".format(dynamic_resource_filename)
                    logger.debug("Pushing dynamic resource to container {:}{:}"
                                 .format(self.container_name, container_path))
                    lxd.file_push(self.container_name, resource_path, container_path)

        num_passes = 0  # Number of test cases passed.
        num_cases = len(test_cases)
        test_case_details = []  # List of dicts each containing the details of a particular test case.

        for i, tc in enumerate(test_cases):
            input_tuple = tc.input_tuple()
            logger.debug("Input tuple: {:}".format(input_tuple))

            if language == 'python':
                runner = PythonRunner()
            elif language == 'javascript':
                runner = JavascriptRunner()
            elif language == 'julia':
                runner = JuliaRunner()
            else:
                raise Exception('Runner not found.')

            try:
                user_answer, process_info = runner.run(self.container_name, code_filename, function_name, input_tuple)

            except (FilePushError, FilePullError):
                explanation = "File could not be pushed to or pulled from LXD container. Returning falcon HTTP 500."
                add_error_to_response(resp, explanation, traceback.format_exc(), falcon.HTTP_500, code_filename)
                return

            except EngineExecutionError:
                explanation = "Return code from executing user code in LXD container is nonzero. " \
                              "Returning falcon HTTP 400."
                add_error_to_response(resp, explanation, traceback.format_exc(), falcon.HTTP_400, code_filename)
                return

            if 'USER_GENERATED_FILES' in tc.input:
                for user_generated_filename in tc.input['USER_GENERATED_FILES']:
                    container_filepath = "/root/{:}".format(user_generated_filename)

                    logger.debug("Pulling user generated file from container {:}{:}"
                                 .format(self.container_name, container_filepath))

                    lxd.file_pull(self.container_name, container_filepath, user_generated_filename)

            if user_answer[0] is None:
                logger.debug("Looks like user's function returned None: output={:}".format(user_answer))
                passed = False
            else:
                try:
                    passed, expected = problem.verify_user_solution(input_tuple, user_answer)
                except Exception:
                    explanation = "Internal engine error during user test case verification. Returning falcon HTTP 500."
                    add_error_to_response(resp, explanation, traceback.format_exc(), falcon.HTTP_500, code_filename)
                    return

            if passed:
                num_passes += 1

            test_case_details.append({
                'testCaseType': tc.test_type.test_name,
                'inputString': str(input_tuple),
                'outputString': str(user_answer),
                'expectedString': str(expected),
                'passed': passed,
                'processInfo': process_info
            })

            if 'DYNAMIC_RESOURCES' in tc.input:
                for dynamic_resource_path in dynamic_resources:
                    logger.debug("Deleting dynamic resource: {:s}".format(dynamic_resource_path))
                    os.remove(dynamic_resource_path)

        logger.info("Passed %d/%d test cases.", num_passes, num_cases)

        resp_dict = {
            'success': True if num_passes == num_cases else False,
            'numTestCases': num_cases,
            'numTestCasesPassed': num_passes,
            'testCaseDetails': test_case_details
        }

        resp.status = falcon.HTTP_200
        resp.set_header('Access-Control-Allow-Origin', '*')
        resp.body = json.dumps(resp_dict)

        util.delete_file(code_filename)
        logger.debug("User code file deleted: {:s}".format(code_filename))

        for file_path in static_resources:
            logging.debug("Deleting static resource {:s}".format(file_path))
            os.remove(file_path)

def parse_payload(http_request):
    try:
        raw_payload_data = http_request.stream.read().decode('utf-8')
    except Exception as ex:
        logger.error("Bad request, reason unknown. Returning 400.")
        raise falcon.HTTPError(falcon.HTTP_400, 'Error', ex.message)

    try:
        json_payload = json.loads(raw_payload_data)
    except ValueError:
        logger.error("Received invalid JSON: {:}".format(raw_payload_data))
        logger.error("Returning 400 error.")
        raise falcon.HTTPError(falcon.HTTP_400, 'Invalid JSON', 'Could not decode request body.')

    return json_payload

def write_code_to_file(code, language):
    """
    Write code into a file with the appropriate file extension.

    :param code: a base64 encoded string representing the user's submitted source code
    :param language: the code's programming language
    :return: the name of the file containing the user's code
    """
    decoded_code = str(base64.b64decode(code), 'utf-8')
    extension = {'python': '.py', 'julia': '.jl', 'javascript': '.js'}.get(language)
    code_filename = util.write_str_to_file(decoded_code, extension)

    logger.debug("User code saved in: {:s}".format(code_filename))

    return code_filename

def add_error_to_response(resp, explanation, tb, falcon_http_error_code, code_filename):
    """
    Modify the falcon HTTP response object with an error to be shown to the user. Also deletes the user's code as the
    engine cannot run it.

    :param resp: The falcon HTTP response object to be modified.
    :param explanation: A human-friendly explanation of the error.
    :param tb: Traceback string.
    :param falcon_http_error_code: Falcon HTTP error code to return.
    :param code_filename: Filepath to user code to be deleted.
    :return: nothing
    """
    logger.error(explanation)
    logger.error(tb)
    util.delete_file(code_filename)

    url_friendly_tb = urllib.parse.quote(tb)  # URL friendly traceback we can embed into a mailto: link.

    DISCOURSE_LINK = 'https://discourse.projectlovelace.net/'
    EMAIL_LINK = 'ada@mg.projectlovelace.net'

    NOTICE = "You should not be seeing this error :( If you have the time, we'd really appreciate\n" \
             "if you could report this on Discourse (" + DISCOURSE_LINK + ") or\n" \
             "via email (" + EMAIL_LINK + "). All the information is embedded in the email link\n" \
             "so all you have to do is press send. Thanks so much!"

    error_message = "{:s}\n\n{:s}\n\nError: {:}".format(explanation, NOTICE, tb)
    resp_dict = {'error': error_message}

    resp.status = falcon_http_error_code
    resp.set_header('Access-Control-Allow-Origin', '*')
    resp.body = json.dumps(resp_dict)
    return

app = falcon.API()
app.add_route('/submit', SubmitResource())
app.add_error_handler(
    Exception,
    lambda ex, req, resp, params: logger.exception(ex)
)
util.configure_lxd()
<commit_msg>Engine API executes all test cases at once<commit_after>import os
import json
import base64
import shutil
import urllib
import logging
import datetime
import importlib
import traceback

import falcon

import engine.util as util
from engine.simple_lxd import simple_lxd as lxd
from engine.runners.python_runner import PythonRunner, FilePushError, FilePullError, EngineExecutionError
from engine.runners.javascript_runner import JavascriptRunner
from engine.runners.julia_runner import JuliaRunner

log_file_path = os.path.join(os.path.dirname(os.path.abspath(__file__)), "logging.ini")
logging.config.fileConfig(log_file_path, disable_existing_loggers=False)
logger = logging.getLogger(__name__)

cwd = os.path.dirname(os.path.abspath(__file__))
os.chdir(cwd)

class SubmitResource(object):
    def __init__(self):
        self.pid = os.getpid()
        self.container_image = "lovelace-image"
        self.container_name = "lovelace-{:d}-{:s}".format(self.pid, datetime.datetime.now().strftime("%Y%m%d-%H%M%S"))

        logger.info("Launching Linux container {:s} (pid={:d}) from image {:s}..."
                    .format(self.container_name, self.pid, self.container_image))

        lxd.launch(self.container_image, name=self.container_name, profile="lovelace")

    def on_post(self, req, resp):
        payload = req.media

        code = payload['code']
        language = payload['language']

        if not code:
            resp_dict = {'error': "No code provided!"}
            resp.status = falcon.HTTP_400
            resp.set_header('Access-Control-Allow-Origin', '*')
            resp.body = json.dumps(resp_dict)
            return

        code_filename = write_code_to_file(code, language)

        try:
            # Fetch problem ID and load the correct problem module.
            problem_name = payload['problem'].replace('-', '_')
            problem_module = 'problems.{:s}'.format(problem_name)

            logger.debug("Importing problem_name={:s} problem_module={:s}...".format(problem_name, problem_module))

            problem = importlib.import_module(problem_module)
        except Exception:
            explanation = "Could not import module {:s}. " \
                          "Returning HTTP 400 Bad Request due to possibly invalid JSON.".format(problem_module)
            add_error_to_response(resp, explanation, traceback.format_exc(), falcon.HTTP_400, code_filename)
            return

        function_name = problem.FUNCTION_NAME
        problem_dir = problem_name

        # Copy static resources into engine directory and push them into the Linux container.
        static_resources = []
        for resource_file_name in problem.STATIC_RESOURCES:
            from_path = os.path.join(cwd, "..", "resources", problem_dir, resource_file_name)
            to_path = os.path.join(cwd, resource_file_name)

            logger.debug("Copying static resource from {:s} to {:s}".format(from_path, to_path))

            try:
                shutil.copyfile(from_path, to_path)
            except Exception:
                explanation = "Engine failed to copy a static resource. Returning falcon HTTP 500."
                add_error_to_response(resp, explanation, traceback.format_exc(), falcon.HTTP_500, code_filename)
                return

            static_resources.append(to_path)

            container_path = "/root/{:}".format(resource_file_name)
            logger.debug("Pushing static resource to container {:}{:}".format(self.container_name, container_path))
            lxd.file_push(self.container_name, from_path, container_path)

        logger.info("Generating test cases...")
        test_cases = []

        try:
            for i, test_type in enumerate(problem.TestCaseType):
                for j in range(test_type.multiplicity):
                    logger.debug("Generating test case {:d}: {:s} ({:d}/{:d})..."
                                 .format(len(test_cases)+1, str(test_type), j+1, test_type.multiplicity))
                    test_cases.append(problem.generate_test_case(test_type))
        except Exception:
            explanation = "Engine failed to generate a test case. Returning falcon HTTP 500."
            add_error_to_response(resp, explanation, traceback.format_exc(), falcon.HTTP_500, code_filename)
            return

        # Copy over all the dynamic resources generated by the test cases.
        dynamic_resources = []
        for i, tc in enumerate(test_cases):
            if 'DYNAMIC_RESOURCES' in tc.input:
                for dynamic_resource_filename in tc.input['DYNAMIC_RESOURCES']:
                    resource_path = os.path.join(cwd, "..", "resources", problem_dir, dynamic_resource_filename)
                    destination_path = os.path.join(cwd, dynamic_resource_filename)

                    logger.debug("Copying test case resource from {:s} to {:s}..."
                                 .format(resource_path, destination_path))

                    shutil.copyfile(resource_path, destination_path)
                    dynamic_resources.append(destination_path)

                    container_path = "/root/{:}".format(dynamic_resource_filename)
                    logger.debug("Pushing dynamic resource to container {:}{:}"
                                 .format(self.container_name, container_path))
                    lxd.file_push(self.container_name, resource_path, container_path)

        if language == 'python':
            runner = PythonRunner()
        elif language == 'javascript':
            runner = JavascriptRunner()
        elif language == 'julia':
            runner = JuliaRunner()
        else:
            raise Exception("Runner not found for language={:}.".format(language))

        try:
            input_tuples = [tc.input() for tc in test_cases]
            user_outputs, p_infos = runner.run(self.container_name, code_filename, function_name, input_tuples)

        except (FilePushError, FilePullError):
            explanation = "File could not be pushed to or pulled from LXD container. Returning falcon HTTP 500."
            add_error_to_response(resp, explanation, traceback.format_exc(), falcon.HTTP_500, code_filename)
            return

        except EngineExecutionError:
            explanation = "Return code from executing user code in LXD container is nonzero. " \
                          "Returning falcon HTTP 400."
            add_error_to_response(resp, explanation, traceback.format_exc(), falcon.HTTP_400, code_filename)
            return

        # Pull any user generated files.
        for i, tc in enumerate(test_cases):
            if 'USER_GENERATED_FILES' in tc.output:
                for user_generated_filename in tc.output['USER_GENERATED_FILES']:
                    container_filepath = "/root/{:s}".format(user_generated_filename)

                    logger.debug("Pulling user generated file from container {:s}{:s}"
                                 .format(self.container_name, container_filepath))

                    lxd.file_pull(self.container_name, container_filepath, user_generated_filename)

        n_cases = len(test_cases)
        n_passes = 0  # Number of test cases passed.
        test_case_details = []  # List of dicts each containing the details of a particular test case.

        # Verify that user outputs are all correct (i.e. check whether each test case passes or fails).
        for input_tuple, user_output, p_info in zip(input_tuples, user_outputs, p_infos):
            if user_output[0] is None:
                logger.debug("Looks like user's function returned None: output={:}".format(user_output))
                passed = False
            else:
                try:
                    passed, expected_output = problem.verify_user_solution(input_tuple, user_output)
                except Exception:
                    explanation = "Internal engine error during user test case verification. Returning falcon HTTP 500."
                    add_error_to_response(resp, explanation, traceback.format_exc(), falcon.HTTP_500, code_filename)
                    return

            if passed:
                n_passes += 1

            test_case_details.append({
                'testCaseType': tc.test_type.test_name,
                'inputString': str(input_tuple),
                'outputString': str(user_output),
                'expectedString': str(expected_output),
                'passed': passed,
                'processInfo': p_info
            })

            if 'DYNAMIC_RESOURCES' in tc.input:
                for dynamic_resource_path in dynamic_resources:
                    logger.debug("Deleting dynamic resource: {:s}".format(dynamic_resource_path))
                    os.remove(dynamic_resource_path)

        logger.info("Passed %d/%d test cases.", n_passes, n_cases)

        resp_dict = {
            'success': True if n_passes == n_cases else False,
            'numTestCases': n_cases,
            'numTestCasesPassed': n_passes,
            'testCaseDetails': test_case_details
        }

        resp.status = falcon.HTTP_200
        resp.set_header('Access-Control-Allow-Origin', '*')
        resp.body = json.dumps(resp_dict)

        util.delete_file(code_filename)
        logger.debug("User code file deleted: {:s}".format(code_filename))

        for file_path in static_resources:
            logging.debug("Deleting static resource {:s}".format(file_path))
            os.remove(file_path)

def parse_payload(http_request):
    try:
        raw_payload_data = http_request.stream.read().decode('utf-8')
    except Exception as ex:
        logger.error("Bad request, reason unknown. Returning 400.")
        raise falcon.HTTPError(falcon.HTTP_400, 'Error', ex.message)

    try:
        json_payload = json.loads(raw_payload_data)
    except ValueError:
        logger.error("Received invalid JSON: {:}".format(raw_payload_data))
        logger.error("Returning 400 error.")
        raise falcon.HTTPError(falcon.HTTP_400, 'Invalid JSON', 'Could not decode request body.')

    return json_payload

def write_code_to_file(code, language):
    """
    Write code into a file with the appropriate file extension.

    :param code: a base64 encoded string representing the user's submitted source code
    :param language: the code's programming language
    :return: the name of the file containing the user's code
    """
    decoded_code = str(base64.b64decode(code), 'utf-8')
    extension = {'python': '.py', 'julia': '.jl', 'javascript': '.js'}.get(language)
    code_filename = util.write_str_to_file(decoded_code, extension)

    logger.debug("User code saved in: {:s}".format(code_filename))

    return code_filename

def add_error_to_response(resp, explanation, tb, falcon_http_error_code, code_filename):
    """
    Modify the falcon HTTP response object with an error to be shown to the user. Also deletes the user's code as the
    engine cannot run it.

    :param resp: The falcon HTTP response object to be modified.
    :param explanation: A human-friendly explanation of the error.
    :param tb: Traceback string.
    :param falcon_http_error_code: Falcon HTTP error code to return.
    :param code_filename: Filepath to user code to be deleted.
    :return: nothing
    """
    logger.error(explanation)
    logger.error(tb)
    util.delete_file(code_filename)

    url_friendly_tb = urllib.parse.quote(tb)  # URL friendly traceback we can embed into a mailto: link.

    DISCOURSE_LINK = 'https://discourse.projectlovelace.net/'
    EMAIL_LINK = 'ada@mg.projectlovelace.net'

    NOTICE = "You should not be seeing this error :( If you have the time, we'd really appreciate\n" \
             "if you could report this on Discourse (" + DISCOURSE_LINK + ") or\n" \
             "via email (" + EMAIL_LINK + "). All the information is embedded in the email link\n" \
             "so all you have to do is press send. Thanks so much!"

    error_message = "{:s}\n\n{:s}\n\nError: {:}".format(explanation, NOTICE, tb)
    resp_dict = {'error': error_message}

    resp.status = falcon_http_error_code
    resp.set_header('Access-Control-Allow-Origin', '*')
    resp.body = json.dumps(resp_dict)
    return

app = falcon.API()
app.add_route('/submit', SubmitResource())
app.add_error_handler(
    Exception,
    lambda ex, req, resp, params: logger.exception(ex)
)
util.configure_lxd()
<|endoftext|>


<commit_before>#!/usr/bin/env python
# encoding: utf-8
"""
http.py

Created by Thomas Mangin on 2011-12-02.
Copyright (c) 2011 Exa Networks. All rights reserved.
"""

#import re
#reg=re.compile('(\w+)[:=] ?"?(\w+)"?')
#reg=re.compile('(\w+)[:=] ?"?([^" ,]+)"?')
#dict(reg.findall(headers))

class Headers (object):
	def __init__ (self,separator):
		self._order = []
		self._data = {}
		self.separator = separator

	def get (self,key,default):
		return self._data.get(key,default)

	def set (self,key,value):
		if key in self._order:
			self._data[key].append(value)
			return self
		self._order.append(key)
		self._data[key] = [value]
		return self

	def replace (self,key,value):
		self._data[key] = [value]

	def default (self,key,value):
		if key not in self._data:
			self._data[key] = [value]

	def extend (self,key,value):
		self._data[key][-1] += value

	def pop (self, key, default=None):
		if key in self._data:
			value = self._data.pop(key)
			self._order.remove(key)
			return value
		else:
			return default

	def parse (self, lines):
		if lines[0].isspace():
			raise ValueError, 'Whitespace before headers'

#		try:
		if True:
			key = ''

			for line in lines.split(self.separator):
				if not line: break

				if line[0].isspace():
					# ValueError if key is not already there
					# IndexError the list is empty
					self.extend(key,line.lstrip())
					continue

				# KeyError if split does not return two elements
				key, value = line.split(':', 1)
				self.set(key.strip().lower(),line)
#		except (KeyError,TypeError,IndexError):
#			raise ValueError, 'Malformed headers'

		# we got a line starting with a :
		if '' in self._data:
			raise ValueError, 'Malformed headers'

		return self

	def __str__ (self):
		return self.separator.join([self.separator.join(self._data[key]) for key in self._order])

<commit_msg>improve header parsing information when it goes wrong<commit_after>#!/usr/bin/env python
# encoding: utf-8
"""
http.py

Created by Thomas Mangin on 2011-12-02.
Copyright (c) 2011 Exa Networks. All rights reserved.
"""

#import re
#reg=re.compile('(\w+)[:=] ?"?(\w+)"?')
#reg=re.compile('(\w+)[:=] ?"?([^" ,]+)"?')
#dict(reg.findall(headers))

class Headers (object):
	def __init__ (self,separator):
		self._order = []
		self._data = {}
		self.separator = separator

	def get (self,key,default):
		return self._data.get(key,default)

	def set (self,key,value):
		if key in self._order:
			self._data[key].append(value)
			return self
		self._order.append(key)
		self._data[key] = [value]
		return self

	def replace (self,key,value):
		self._data[key] = [value]

	def default (self,key,value):
		if key not in self._data:
			self._data[key] = [value]

	def extend (self,key,value):
		self._data[key][-1] += value

	def pop (self, key, default=None):
		if key in self._data:
			value = self._data.pop(key)
			self._order.remove(key)
			return value
		else:
			return default

	def parse (self, lines):
		if lines[0].isspace():
			raise ValueError('Malformed headers, headers starts with a white space')

		try:
			key = ''

			for line in lines.split(self.separator):
				if not line: break

				if line[0].isspace():
					# ValueError if key is not already there
					# IndexError the list is empty
					self.extend(key,line.lstrip())
					continue

				# KeyError if split does not return two elements
				key, value = line.split(':', 1)
				self.set(key.strip().lower(),line)
		except (KeyError,TypeError,IndexError):
			raise ValueError('Malformed headers (line : %s) headers %s' % (line,lines.replace('\n','\\n').replace('\r','\\r')))

		# we got a line starting with a :
		if '' in self._data:
			raise ValueError ('Malformed headers, line starts with colon (:)')

		return self

	def __str__ (self):
		return self.separator.join([self.separator.join(self._data[key]) for key in self._order])

<|endoftext|>