<commit_before># encoding: utf-8
""" Contains code for parsing datatypes from the statscraper-datatypes repo
"""

from glob import iglob
from itertools import chain
from csvkit import DictReader
from csvkit import reader as CsvReader
from . import base_scraper  # FIXME move DimensionValue to avoid this import loop
import os

DIR_PATH = os.path.dirname(os.path.realpath(__file__))
DATATYPES_FILE = os.path.join(DIR_PATH, "datatypes", "datatypes.csv")
VALUE_DELIMITOR = ','

class NoSuchDatatype(Exception):
    """No datatype with that id."""

    pass

class Datatype(object):
    """Represent a datatype, initiated by id."""

    def __init__(self, id):
        """Id is a datatype from datatypes.csv."""
        self.id = id
        self.allowed_values = base_scraper.ValueList()

        data = None
        with open(DATATYPES_FILE, 'rb') as csvfile:
            reader = DictReader(csvfile)
            for row in reader:
                if row["id"] == id:
                    data = row
                    break
        if data is None:
            raise(NoSuchDatatype)
        self.value_type = data["value_type"]
        self.description = data["description"]
        domain = data["allowed_values"]
        if domain:
            for file_ in self._get_csv_files(domain):
                with open(file_, 'rb') as csvfile:
                    reader = DictReader(csvfile)
                    dialect_names = [x
                                     for x in reader.fieldnames
                                     if x.startswith("dialect:")]
                    self.dialects = [d[8:] for d in dialect_names]
                    for row in reader:
                        value = base_scraper.DimensionValue(row["id"],
                                                            self,
                                                            label=row["label"])
                        dialects = {x: None for x in self.dialects}
                        from StringIO import StringIO
                        for d in dialect_names:
                            f = StringIO(row[d])
                            csvreader = CsvReader(f, delimiter=VALUE_DELIMITOR,
                                                  skipinitialspace=True,
                                                  strict=True)
                            try:
                                values = csvreader.next()
                            except Exception:
                                continue
                            dialects[d[8:]] = values
                        value.dialects = dialects
                        self.allowed_values.append(value)

    def _get_csv_files(self, domain):
        domain = os.path.join(*domain.split("/"))

        # We are fetching both by filename and dir name
        # so that regions/kenya will match anything in
        # `datatypes/values/regions/kenya/*.csv`
        # and/or `datatypes/values/regions/kenya.csv`
        #
        # There is probably an easier way to do this
        # FIXME the below function fetches /foo/bar/regions/kenya as well, but we probably want ^regions/kenya
        value_path_1 = os.path.join(DIR_PATH, "datatypes", "values", domain)
        value_path_2 = os.path.join(DIR_PATH, "datatypes", "values")
        files_1 = chain.from_iterable(iglob(os.path.join(root, '*.csv'))
                                      for root, dirs, files in os.walk(value_path_1))
        files_2 = chain.from_iterable(iglob(os.path.join(root, domain + '.csv'))
                                      for root, dirs, files in os.walk(value_path_2))
        for f in chain(files_1, files_2):
            yield f

    def __str__(self):
        return str(self.id)

    def __repr__(self):
        return '<Datatype: %s>' % str(self)
<commit_msg>compatible stringio import<commit_after># encoding: utf-8
""" Contains code for parsing datatypes from the statscraper-datatypes repo
"""

from glob import iglob
from itertools import chain
from csvkit import DictReader
from csvkit import reader as CsvReader
from . import base_scraper  # FIXME move DimensionValue to avoid this import loop
import os

DIR_PATH = os.path.dirname(os.path.realpath(__file__))
DATATYPES_FILE = os.path.join(DIR_PATH, "datatypes", "datatypes.csv")
VALUE_DELIMITOR = ','

class NoSuchDatatype(Exception):
    """No datatype with that id."""

    pass

class Datatype(object):
    """Represent a datatype, initiated by id."""

    def __init__(self, id):
        """Id is a datatype from datatypes.csv."""
        self.id = id
        self.allowed_values = base_scraper.ValueList()

        data = None
        with open(DATATYPES_FILE, 'rb') as csvfile:
            reader = DictReader(csvfile)
            for row in reader:
                if row["id"] == id:
                    data = row
                    break
        if data is None:
            raise(NoSuchDatatype)
        self.value_type = data["value_type"]
        self.description = data["description"]
        domain = data["allowed_values"]
        if domain:
            for file_ in self._get_csv_files(domain):
                with open(file_, 'rb') as csvfile:
                    reader = DictReader(csvfile)
                    dialect_names = [x
                                     for x in reader.fieldnames
                                     if x.startswith("dialect:")]
                    self.dialects = [d[8:] for d in dialect_names]
                    for row in reader:
                        value = base_scraper.DimensionValue(row["id"],
                                                            self,
                                                            label=row["label"])
                        dialects = {x: None for x in self.dialects}

                        # TODO: Shouldn't this be moved to the top?
                        try:
                            from StringIO import StringIO
                        except ImportError:
                            from io import StringIO

                        for d in dialect_names:
                            f = StringIO(row[d])
                            csvreader = CsvReader(f, delimiter=VALUE_DELIMITOR,
                                                  skipinitialspace=True,
                                                  strict=True)
                            try:
                                values = csvreader.next()
                            except Exception:
                                continue
                            dialects[d[8:]] = values
                        value.dialects = dialects
                        self.allowed_values.append(value)

    def _get_csv_files(self, domain):
        domain = os.path.join(*domain.split("/"))

        # We are fetching both by filename and dir name
        # so that regions/kenya will match anything in
        # `datatypes/values/regions/kenya/*.csv`
        # and/or `datatypes/values/regions/kenya.csv`
        #
        # There is probably an easier way to do this
        # FIXME the below function fetches /foo/bar/regions/kenya as well, but we probably want ^regions/kenya
        value_path_1 = os.path.join(DIR_PATH, "datatypes", "values", domain)
        value_path_2 = os.path.join(DIR_PATH, "datatypes", "values")
        files_1 = chain.from_iterable(iglob(os.path.join(root, '*.csv'))
                                      for root, dirs, files in os.walk(value_path_1))
        files_2 = chain.from_iterable(iglob(os.path.join(root, domain + '.csv'))
                                      for root, dirs, files in os.walk(value_path_2))
        for f in chain(files_1, files_2):
            yield f

    def __str__(self):
        return str(self.id)

    def __repr__(self):
        return '<Datatype: %s>' % str(self)
<|endoftext|>


<commit_before>import SocketServer
import parser
import metrics
import logging

class CollectorHandler(SocketServer.RequestHandler):
    """Simple handler that parses a UDP packet and adds the metrics to an aggregator"""
    def handle(self):
        # Get the raw metrics
        mesg, socket = self.request
        raw_metrics = parser.parse_message(mesg)

        # Convert to the metrics objects
        metric_lst = []
        for (key,value,metric_type,flag) in raw_metrics:
            if metric_type in metrics.METRIC_TYPES:
                # Create and store the metric object
                metric = metrics.METRIC_TYPE[metric_type](key,value,flag)
                metric_lst.append(metric)
            else:
                self.server.logger.error("Invalid metric type '%s' received! Dropping." % metric_type)

        # Add all the metrics to the aggregator
        self.server.aggregator.add_metrics(metrics)

class CollectorServer(SocketServer.UDPServer):
    """UDP listener that invokes the CollectorHandler for each packet"""
    allow_reuse_address = True

    def __init__(self,hostname,port,aggregator,handler=CollectorHandler):
        SocketServer.UDPServer.__init__(self,(hostname,port),handler)
        self.aggregator = aggregator
        self.logger = logging.getLogger("statsite.udpcollector")

<commit_msg>Remove udpcollector.py<commit_after><|endoftext|>


<commit_before>
from webdoc import sequence

class op(object):
	def __repr__(self):
		return self.__class__.__name__
	def __eq__(self, x):
		return isinstance(x, self.__class__)

class EQUAL(op): pass
EQUAL = EQUAL()
class NOTEQUAL(op): pass
NOTEQUAL = NOTEQUAL()
class LESSTHAN(op): pass
LESSTHAN = LESSTHAN()
class LESSEQUAL(op): pass
LESSEQUAL = LESSEQUAL()
class GREATERTHAN(op): pass
GREATERTHAN = GREATERTHAN()
class GREATEREQUAL(op): pass
GREATEREQUAL = GREATEREQUAL()
class ADD(op): pass
ADD = ADD()
class CONCATENATE(op): pass
CONCATENATE = CONCATENATE()
class SUBTRACT(op): pass
SUBTRACT = SUBTRACT()
class MULTIPLY(op): pass
MULTIPLY = MULTIPLY()
class DIVIDE(op): pass
DIVIDE = DIVIDE()
class FLOORDIVIDE(op): pass
FLOORDIVIDE = FLOORDIVIDE()
class MODULO(op): pass
MODULO = MODULO()
class AND(op): pass
AND = AND()
class OR(op): pass
OR = OR()
class NOT(op): pass
NOT = NOT()
class NEGATIVE(op): pass
NEGATIVE = NEGATIVE()
class ABS(op): pass
ABS = ABS()
class LENGTH(op): pass
LENGTH = LENGTH()
class ASCEND(op): pass
ASCEND = ASCEND()
class DESCEND(op): pass
DESCEND = DESCEND()
class SUM(op): pass
SUM = SUM()
class AVERAGE(op): pass
AVERAGE = AVERAGE()
class BETWEEN(op): pass
BETWEEN = BETWEEN()
class MIN(op): pass
MIN = MIN()
class MAX(op): pass
MAX = MAX()
class UPPER(op): pass
UPPER = UPPER()
class LOWER(op): pass
LOWER = LOWER()
class LIKE(op): pass
LIKE = LIKE()
class GLOB(op): pass
GLOB = GLOB()
class LSTRIP(op): pass
LSTRIP = LSTRIP()
class STRIP(op): pass
STRIP = STRIP()
class RSTRIP(op): pass
RSTRIP = RSTRIP()
class REPLACE(op): pass
REPLACE = REPLACE()
class ROUND(op): pass
ROUND = ROUND()
class SUBSTRING(op): pass
SUBSTRING = SUBSTRING()
class COALESCE(op): pass
COALESCE = COALESCE()

class driver_base(object):
	'''Base class for database drivers

	This class abstracts away a lot of the logic needed to implement a database
	driver for webdb. Derived classes must overload the following methods to
	have a working driver. If the task is best accomplished using a single SQL
	command, the method ending in '_sql' should be defined, as the alternatives
	use those to accomplish their tasks. For more information, view documentation
	on each specific method

	  * list_tables or list_tables_sql
	     - list all tables in the database. Implements: db.conform
	  * list_columns or list_columns_sql
	     - list all columns defined in a table. Implements: db.conform
	  * create_table, create_table_sql, create_table_if_nexists,
	    or create_table_if_nexists_sql
	     - create tables (if missing). Implements: db.__setattr__, db.migrate
	  * rename_table or rename_table_sql.
	     - changes a table's name. Implements: db.migrate
	  * add_column or add_column_sql
	     - adds a new column to a table. Implements: db.migrate
	  * select or select_sql
	     - retrieves rows from a table. Implements: Where.select, table.select
	  * insert or insert_sql
	     - adds rows to a table. Implements: table.insert
	  * update or update_sql
	     - alters records in a table. Implements: Where.update, row.update
	  * delete or delete_sql
	     - removes records from a table. Implements: Where.delete, row.delete,
	       table.__delitem__, table.__delattr__

	Additionally, the following variables must be defined.

	  * webdb_types: a dictionary mapping webdb column types to names of database
	      column types. Used for defining tables
	  * driver_types: a dictionary mapping database column types to webdb column
	      types. Used when conforming.

	The following methods should be defined by subclasses if the database uses
	non-standard syntax

	  * identifier
	     - checks that a table or column name uses valid characters, and is properly
	       escaped (to avoid keywords). Default encloses name in double quotes (")
	  * literal
	     - formats a literal value to be used in an sql expression.
	  * column_name
	     - returns the name of a column. Default returns dot (.) joined result
	       of identifier on its arguments, i.e. "table"."column"
	  * format_column
	     - returns a column definition for its Column object argument

	The following methods may be defined by subclasses, but are not required for
	normal use.

	  * drop_column or drop_column_sql
	     - removes a column and all its data from a table. Columns in a table
	       which don't appear in a table definition are ignored.
	       Implements: table.drop_column
	'''

	def __init__(self, connection):
		self.connection = connection
		self.depth = 0
		self.cursor = None

	def __enter__(self):
		self.depth += 1
		if self.cursor is None:
			self.cursor = self.connection.cursor()
		return self.cursor

	def __exit__(self, obj, exc, tb):
		self.depth -= 1
		if self.depth == 0:
			if obj:
				self.connection.rollback()
			else:
				self.connection.commit()
			self.cursor = None

	def execute(self, sql, values=()):
		self.lastsql = sql
		with self as cursor:
			try:
				return cursor.execute(sql, values)
			except Exception, e:
				self.handle_exception(e)
				raise Exception(e, sql, values)

	def identifier(self, name):
		if not name.replace('_','').isalnum():
			raise NameError("Column names can only contain letters, numbers, and underscores. Got %r" % name)
		return '"%s"'%name

	def literal(self, value, cast=None):
		if value is None:
			return 'NULL'
		elif isinstance(value, basestring) or cast in ('TEXT','BLOB'):
			return "'%s'"%str(value).replace("'", "''")
		elif cast in ('INT', 'REAL'):
			return '%g'%value
		else:
			return value

	def column_name(self, table, col):
		return '.'.join(map(self.identifier, (table, col)))

	def expression(self, x):
		if isinstance(x, list):
			operator = x[0]
			return '(%s)'%getattr(self,'op_%s'%operator)(*map(self.expression,x[1:]))
		elif hasattr(x, 'table') and hasattr(x, 'name'): #Column duck-typed
			return self.column_name(x.table._name, x.name)
		elif hasattr(x, '_where_tree'): #Where duck-typed
			return self.expression(x._where_tree)
		else:
			return self.literal(x)

	def parse_where(self, where_clause):
		if where_clause:
			clause = self.expression(where_clause)
			if clause:
				clause = ' WHERE '+clause
		else:
			clause = ''
		return clause

	def format_column(self, column):
		if hasattr(column,'reftable'):
			type = self.map_type(column.type) % dict(table=column.reftable._name, column=column.reftable.rowid.name)
		else:
			type = self.map_type(column.type)
		default = " DEFAULT %s"%self.literal(column.default, type) if not callable(column.default) and (column.notnull or not column.default is None) else ''
		return '%(name)s %(type)s%(notnull)s%(default)s' % {
			'name': self.identifier(column.name),
			'type': type,
			'notnull': ' NOT NULL' if column.notnull else '',
			'default': default
		}

	def map_type(self, t):
		r = self.webdb_types.get(t)
		if r:
			return r
		else:
			raise Exception('Unknown column type %s' % t)

	def unmap_type(self, t):
		r = self.driver_types.get(t)
		if r:
			return r
		else:
			raise Exception('Unknown column type %s' % t)

	def create_table_if_nexists_sql(self, name, *columns):
		'''create_table_if_nexists_sql(self, name, *columns) -> Stub

		Base classes should return SQL for creating a table if it doesn't
		exist. '''
		raise NotImplementedError

	def create_table_if_nexists(self, name, table):
		if hasattr(self, 'create_table_if_nexists_sql'):
			self.execute(self.create_table_if_nexists_sql(
				self.identifier(name),
				*map(self.format_column, table.columns)
			))
		elif name not in self.list_tables:
			self.create_table(name, table.columns)

	def create_table_sql(self, name, *columns):
		raise NotImplementedError

	def create_table(self, name, table):
		try:
			self.execute(self.create_table_sql(self.identifier(name), *map(self.format_column,table.columns)))
		except NotImplementedError:
			self.create_table_if_nexists(name, table.columns)

	def list_tables(self):
		return (str(n) for (n,) in self.execute(self.list_tables_sql()))

	def rename_table(self, orig, new):
		self.execute(self.rename_table_sql(self.identifier(orig), self.identifier(new)))

	def alter_table(self, name, table):
		with self:
			db_cols = self.list_columns(name)
			db_names = [c[0] for c in db_cols]
			for column in table.columns:
				if column.name not in db_names:
					self.execute(self.add_column_sql(self.identifier(name), self.format_column(column)))

	def select(self, columns, tables, conditions, props):
		return self.execute(self.select_sql(
			([self.identifier('rowid')]if len(tables)==1 and not props.get('distinct') else [])+map(self.expression,columns),
			[self.identifier(t._name) for t in tables],
			self.parse_where(conditions),
			props.get('distinct',False),
			sequence(props.get('orderby',[])),
		))

	def insert(self, table, values):
		cur = self.execute(self.insert_sql(self.identifier(table), map(self.identifier,values.keys())), values.values())
		return cur.lastrowid

	def update(self, table, conditions, values):
		return self.execute(self.update_sql(self.identifier(table), map(self.identifier,values.keys()), self.parse_where(conditions)), values.values())

	def delete(self, table, conditions):
		return self.execute(self.delete_sql(self.identifier(table), self.parse_where(conditions)))

	op_EQUAL = staticmethod(lambda a,b:'%s=%s'%(a,b))
	op_LESSEQUAL = staticmethod(lambda a,b:'%s<=%s'%(a,b))
	op_GREATERTHAN = staticmethod(lambda a,b:'%s>%s'%(a,b))
	op_NOTEQUAL = staticmethod(lambda a,b:'%s!=%s'%(a,b))
	op_LESSTHAN = staticmethod(lambda a,b:'%s<%s'%(a,b))
	op_GREATEREQUAL = staticmethod(lambda a,b:'%s>=%s'%(a,b))
	op_ADD = staticmethod(lambda a,b:'%s+%s'%(a,b))
	op_CONCATENATE = staticmethod(lambda a,b:'%s||%s'%(a,b))
	op_SUBTRACT = staticmethod(lambda a,b:'%s-%s'%(a,b))
	op_MULTIPLY = staticmethod(lambda a,b:'%s*%s'%(a,b))
	op_DIVIDE = staticmethod(lambda a,b:'%s/%s'%(a,b))
	op_FLOORDIVIDE = staticmethod(lambda a,b:'%s/%s'%(a,b))
	op_MODULO = staticmethod(lambda a,b:'%s%%%s'%(a,b))
	op_AND = staticmethod(lambda a,b:'%s AND %s'%(a,b))
	op_OR = staticmethod(lambda a,b:'%s OR %s'%(a,b))
	op_NOT = staticmethod(lambda a:'NOT %s'%a)
	op_NEGATIVE = staticmethod(lambda a:'-%s'%a)
	op_ABS = staticmethod(lambda a:'abs(%s)'%a)
	op_LENGTH = staticmethod(lambda a:'length(%s)'%a)
	op_ASCEND = staticmethod(lambda a:'%s ASC'%a)
	op_DESCEND = staticmethod(lambda a:'%s DESC'%a)
	op_SUM = staticmethod(lambda a:'total(%s)'%a)
	op_AVERAGE = staticmethod(lambda a:'avg(%s)'%a)
	op_BETWEEN = staticmethod(lambda a,b,c:'%s BETWEEN %s AND %s'%(a,b,c))
	op_MIN = staticmethod(lambda a:'min(%s)'%a)
	op_MAX = staticmethod(lambda a:'max(%s)'%a)
	op_UPPER = staticmethod(lambda a:'upper(%s)'%a)
	op_LOWER = staticmethod(lambda a:'lower(%s)'%a)
	op_LIKE = staticmethod(lambda a,b,c=None:'%s LIKE %s'%(a,b) if c is None else '%s LIKE %s ESCAPE %s'%(a,b,c))
	op_SUBSTRING = staticmethod(lambda a,b,c=None:'substr(%s,%s)'%(a,b) if c is None else 'substr(%s,%s,%s)'%(a,b,c))
	op_GLOB = staticmethod(lambda a,b:'%s GLOB %s' % (a,b))
	op_LSTRIP = staticmethod(lambda a,b=None:'ltrim(%s)'%a if b is None else 'ltrim(%s,%s)'%(a,b))
	op_RSTRIP = staticmethod(lambda a,b=None:'rtrim(%s)'%a if b is None else 'rtrim(%s,%s)'%(a,b))
	op_STRIP = staticmethod(lambda a,b=None:'trim(%s)'%a if b is None else 'trim(%s,%s)'%(a,b))
	op_REPLACE = staticmethod(lambda a,b,c:'replace(%s,%s,%s)'%(a,b,c))
	op_ROUND = staticmethod(lambda a,b=None:'round(%s)'%a if b is None else 'round(%s,%s'%(a,b))
	op_COALESCE = staticmethod(lambda a,b,*c:'coalesce(%s,%s,%s)'%(a,b,','.join(c)) if c else 'coalesce(%s,%s)'%(a,b))
<commit_msg>Fixed old import<commit_after>
from silk import sequence

class op(object):
	def __repr__(self):
		return self.__class__.__name__
	def __eq__(self, x):
		return isinstance(x, self.__class__)

class EQUAL(op): pass
EQUAL = EQUAL()
class NOTEQUAL(op): pass
NOTEQUAL = NOTEQUAL()
class LESSTHAN(op): pass
LESSTHAN = LESSTHAN()
class LESSEQUAL(op): pass
LESSEQUAL = LESSEQUAL()
class GREATERTHAN(op): pass
GREATERTHAN = GREATERTHAN()
class GREATEREQUAL(op): pass
GREATEREQUAL = GREATEREQUAL()
class ADD(op): pass
ADD = ADD()
class CONCATENATE(op): pass
CONCATENATE = CONCATENATE()
class SUBTRACT(op): pass
SUBTRACT = SUBTRACT()
class MULTIPLY(op): pass
MULTIPLY = MULTIPLY()
class DIVIDE(op): pass
DIVIDE = DIVIDE()
class FLOORDIVIDE(op): pass
FLOORDIVIDE = FLOORDIVIDE()
class MODULO(op): pass
MODULO = MODULO()
class AND(op): pass
AND = AND()
class OR(op): pass
OR = OR()
class NOT(op): pass
NOT = NOT()
class NEGATIVE(op): pass
NEGATIVE = NEGATIVE()
class ABS(op): pass
ABS = ABS()
class LENGTH(op): pass
LENGTH = LENGTH()
class ASCEND(op): pass
ASCEND = ASCEND()
class DESCEND(op): pass
DESCEND = DESCEND()
class SUM(op): pass
SUM = SUM()
class AVERAGE(op): pass
AVERAGE = AVERAGE()
class BETWEEN(op): pass
BETWEEN = BETWEEN()
class MIN(op): pass
MIN = MIN()
class MAX(op): pass
MAX = MAX()
class UPPER(op): pass
UPPER = UPPER()
class LOWER(op): pass
LOWER = LOWER()
class LIKE(op): pass
LIKE = LIKE()
class GLOB(op): pass
GLOB = GLOB()
class LSTRIP(op): pass
LSTRIP = LSTRIP()
class STRIP(op): pass
STRIP = STRIP()
class RSTRIP(op): pass
RSTRIP = RSTRIP()
class REPLACE(op): pass
REPLACE = REPLACE()
class ROUND(op): pass
ROUND = ROUND()
class SUBSTRING(op): pass
SUBSTRING = SUBSTRING()
class COALESCE(op): pass
COALESCE = COALESCE()

class driver_base(object):
	'''Base class for database drivers

	This class abstracts away a lot of the logic needed to implement a database
	driver for webdb. Derived classes must overload the following methods to
	have a working driver. If the task is best accomplished using a single SQL
	command, the method ending in '_sql' should be defined, as the alternatives
	use those to accomplish their tasks. For more information, view documentation
	on each specific method

	  * list_tables or list_tables_sql
	     - list all tables in the database. Implements: db.conform
	  * list_columns or list_columns_sql
	     - list all columns defined in a table. Implements: db.conform
	  * create_table, create_table_sql, create_table_if_nexists,
	    or create_table_if_nexists_sql
	     - create tables (if missing). Implements: db.__setattr__, db.migrate
	  * rename_table or rename_table_sql.
	     - changes a table's name. Implements: db.migrate
	  * add_column or add_column_sql
	     - adds a new column to a table. Implements: db.migrate
	  * select or select_sql
	     - retrieves rows from a table. Implements: Where.select, table.select
	  * insert or insert_sql
	     - adds rows to a table. Implements: table.insert
	  * update or update_sql
	     - alters records in a table. Implements: Where.update, row.update
	  * delete or delete_sql
	     - removes records from a table. Implements: Where.delete, row.delete,
	       table.__delitem__, table.__delattr__

	Additionally, the following variables must be defined.

	  * webdb_types: a dictionary mapping webdb column types to names of database
	      column types. Used for defining tables
	  * driver_types: a dictionary mapping database column types to webdb column
	      types. Used when conforming.

	The following methods should be defined by subclasses if the database uses
	non-standard syntax

	  * identifier
	     - checks that a table or column name uses valid characters, and is properly
	       escaped (to avoid keywords). Default encloses name in double quotes (")
	  * literal
	     - formats a literal value to be used in an sql expression.
	  * column_name
	     - returns the name of a column. Default returns dot (.) joined result
	       of identifier on its arguments, i.e. "table"."column"
	  * format_column
	     - returns a column definition for its Column object argument

	The following methods may be defined by subclasses, but are not required for
	normal use.

	  * drop_column or drop_column_sql
	     - removes a column and all its data from a table. Columns in a table
	       which don't appear in a table definition are ignored.
	       Implements: table.drop_column
	'''

	def __init__(self, connection):
		self.connection = connection
		self.depth = 0
		self.cursor = None

	def __enter__(self):
		self.depth += 1
		if self.cursor is None:
			self.cursor = self.connection.cursor()
		return self.cursor

	def __exit__(self, obj, exc, tb):
		self.depth -= 1
		if self.depth == 0:
			if obj:
				self.connection.rollback()
			else:
				self.connection.commit()
			self.cursor = None

	def execute(self, sql, values=()):
		self.lastsql = sql
		with self as cursor:
			try:
				return cursor.execute(sql, values)
			except Exception, e:
				self.handle_exception(e)
				raise Exception(e, sql, values)

	def identifier(self, name):
		if not name.replace('_','').isalnum():
			raise NameError("Column names can only contain letters, numbers, and underscores. Got %r" % name)
		return '"%s"'%name

	def literal(self, value, cast=None):
		if value is None:
			return 'NULL'
		elif isinstance(value, basestring) or cast in ('TEXT','BLOB'):
			return "'%s'"%str(value).replace("'", "''")
		elif cast in ('INT', 'REAL'):
			return '%g'%value
		else:
			return value

	def column_name(self, table, col):
		return '.'.join(map(self.identifier, (table, col)))

	def expression(self, x):
		if isinstance(x, list):
			operator = x[0]
			return '(%s)'%getattr(self,'op_%s'%operator)(*map(self.expression,x[1:]))
		elif hasattr(x, 'table') and hasattr(x, 'name'): #Column duck-typed
			return self.column_name(x.table._name, x.name)
		elif hasattr(x, '_where_tree'): #Where duck-typed
			return self.expression(x._where_tree)
		else:
			return self.literal(x)

	def parse_where(self, where_clause):
		if where_clause:
			clause = self.expression(where_clause)
			if clause:
				clause = ' WHERE '+clause
		else:
			clause = ''
		return clause

	def format_column(self, column):
		if hasattr(column,'reftable'):
			type = self.map_type(column.type) % dict(table=column.reftable._name, column=column.reftable.rowid.name)
		else:
			type = self.map_type(column.type)
		default = " DEFAULT %s"%self.literal(column.default, type) if not callable(column.default) and (column.notnull or not column.default is None) else ''
		return '%(name)s %(type)s%(notnull)s%(default)s' % {
			'name': self.identifier(column.name),
			'type': type,
			'notnull': ' NOT NULL' if column.notnull else '',
			'default': default
		}

	def map_type(self, t):
		r = self.webdb_types.get(t)
		if r:
			return r
		else:
			raise Exception('Unknown column type %s' % t)

	def unmap_type(self, t):
		r = self.driver_types.get(t)
		if r:
			return r
		else:
			raise Exception('Unknown column type %s' % t)

	def create_table_if_nexists_sql(self, name, *columns):
		'''create_table_if_nexists_sql(self, name, *columns) -> Stub

		Base classes should return SQL for creating a table if it doesn't
		exist. '''
		raise NotImplementedError

	def create_table_if_nexists(self, name, table):
		if hasattr(self, 'create_table_if_nexists_sql'):
			self.execute(self.create_table_if_nexists_sql(
				self.identifier(name),
				*map(self.format_column, table.columns)
			))
		elif name not in self.list_tables:
			self.create_table(name, table.columns)

	def create_table_sql(self, name, *columns):
		raise NotImplementedError

	def create_table(self, name, table):
		try:
			self.execute(self.create_table_sql(self.identifier(name), *map(self.format_column,table.columns)))
		except NotImplementedError:
			self.create_table_if_nexists(name, table.columns)

	def list_tables(self):
		return (str(n) for (n,) in self.execute(self.list_tables_sql()))

	def rename_table(self, orig, new):
		self.execute(self.rename_table_sql(self.identifier(orig), self.identifier(new)))

	def alter_table(self, name, table):
		with self:
			db_cols = self.list_columns(name)
			db_names = [c[0] for c in db_cols]
			for column in table.columns:
				if column.name not in db_names:
					self.execute(self.add_column_sql(self.identifier(name), self.format_column(column)))

	def select(self, columns, tables, conditions, props):
		return self.execute(self.select_sql(
			([self.identifier('rowid')]if len(tables)==1 and not props.get('distinct') else [])+map(self.expression,columns),
			[self.identifier(t._name) for t in tables],
			self.parse_where(conditions),
			props.get('distinct',False),
			sequence(props.get('orderby',[])),
		))

	def insert(self, table, values):
		cur = self.execute(self.insert_sql(self.identifier(table), map(self.identifier,values.keys())), values.values())
		return cur.lastrowid

	def update(self, table, conditions, values):
		return self.execute(self.update_sql(self.identifier(table), map(self.identifier,values.keys()), self.parse_where(conditions)), values.values())

	def delete(self, table, conditions):
		return self.execute(self.delete_sql(self.identifier(table), self.parse_where(conditions)))

	op_EQUAL = staticmethod(lambda a,b:'%s=%s'%(a,b))
	op_LESSEQUAL = staticmethod(lambda a,b:'%s<=%s'%(a,b))
	op_GREATERTHAN = staticmethod(lambda a,b:'%s>%s'%(a,b))
	op_NOTEQUAL = staticmethod(lambda a,b:'%s!=%s'%(a,b))
	op_LESSTHAN = staticmethod(lambda a,b:'%s<%s'%(a,b))
	op_GREATEREQUAL = staticmethod(lambda a,b:'%s>=%s'%(a,b))
	op_ADD = staticmethod(lambda a,b:'%s+%s'%(a,b))
	op_CONCATENATE = staticmethod(lambda a,b:'%s||%s'%(a,b))
	op_SUBTRACT = staticmethod(lambda a,b:'%s-%s'%(a,b))
	op_MULTIPLY = staticmethod(lambda a,b:'%s*%s'%(a,b))
	op_DIVIDE = staticmethod(lambda a,b:'%s/%s'%(a,b))
	op_FLOORDIVIDE = staticmethod(lambda a,b:'%s/%s'%(a,b))
	op_MODULO = staticmethod(lambda a,b:'%s%%%s'%(a,b))
	op_AND = staticmethod(lambda a,b:'%s AND %s'%(a,b))
	op_OR = staticmethod(lambda a,b:'%s OR %s'%(a,b))
	op_NOT = staticmethod(lambda a:'NOT %s'%a)
	op_NEGATIVE = staticmethod(lambda a:'-%s'%a)
	op_ABS = staticmethod(lambda a:'abs(%s)'%a)
	op_LENGTH = staticmethod(lambda a:'length(%s)'%a)
	op_ASCEND = staticmethod(lambda a:'%s ASC'%a)
	op_DESCEND = staticmethod(lambda a:'%s DESC'%a)
	op_SUM = staticmethod(lambda a:'total(%s)'%a)
	op_AVERAGE = staticmethod(lambda a:'avg(%s)'%a)
	op_BETWEEN = staticmethod(lambda a,b,c:'%s BETWEEN %s AND %s'%(a,b,c))
	op_MIN = staticmethod(lambda a:'min(%s)'%a)
	op_MAX = staticmethod(lambda a:'max(%s)'%a)
	op_UPPER = staticmethod(lambda a:'upper(%s)'%a)
	op_LOWER = staticmethod(lambda a:'lower(%s)'%a)
	op_LIKE = staticmethod(lambda a,b,c=None:'%s LIKE %s'%(a,b) if c is None else '%s LIKE %s ESCAPE %s'%(a,b,c))
	op_SUBSTRING = staticmethod(lambda a,b,c=None:'substr(%s,%s)'%(a,b) if c is None else 'substr(%s,%s,%s)'%(a,b,c))
	op_GLOB = staticmethod(lambda a,b:'%s GLOB %s' % (a,b))
	op_LSTRIP = staticmethod(lambda a,b=None:'ltrim(%s)'%a if b is None else 'ltrim(%s,%s)'%(a,b))
	op_RSTRIP = staticmethod(lambda a,b=None:'rtrim(%s)'%a if b is None else 'rtrim(%s,%s)'%(a,b))
	op_STRIP = staticmethod(lambda a,b=None:'trim(%s)'%a if b is None else 'trim(%s,%s)'%(a,b))
	op_REPLACE = staticmethod(lambda a,b,c:'replace(%s,%s,%s)'%(a,b,c))
	op_ROUND = staticmethod(lambda a,b=None:'round(%s)'%a if b is None else 'round(%s,%s'%(a,b))
	op_COALESCE = staticmethod(lambda a,b,*c:'coalesce(%s,%s,%s)'%(a,b,','.join(c)) if c else 'coalesce(%s,%s)'%(a,b))
<|endoftext|>


<commit_before>import asyncio
from aiohttp import web

from concurrent.futures import ThreadPoolExecutor
import threading
from time import sleep

async def handle(request):
    index = open("index.html", 'rb')
    content = index.read()
    return web.Response(body=content)

tick = asyncio.Condition()

async def wshandler(request):
    ws = web.WebSocketResponse()
    await ws.prepare(request)

    recv_task = None
    tick_task = None
    while 1:
        if not recv_task:
            recv_task = asyncio.ensure_future(ws.receive())
        if not tick_task:
            await tick.acquire()
            tick_task = asyncio.ensure_future(tick.wait())

        done, pending = await asyncio.wait(
            [recv_task,
             tick_task],
            return_when=asyncio.FIRST_COMPLETED)

        if recv_task in done:
            msg = recv_task.result()
            if msg.tp == web.MsgType.text:
                print("Got message %s" % msg.data)
                ws.send_str("Pressed key code: {}".format(msg.data))
            elif msg.tp == web.MsgType.close or\
                 msg.tp == web.MsgType.error:
                break
            recv_task = None

        if tick_task in done:
            ws.send_str("game loop ticks")
            tick.release()
            tick_task = None

    return ws

def game_loop(asyncio_loop):
    print("Game loop thread id {}".format(threading.get_ident()))
    # a coroutine to run in main thread
    async def notify():
        print("Notify thread id {}".format(threading.get_ident()))
        await tick.acquire()
        tick.notify_all()
        tick.release()

    while 1:
        task = asyncio.run_coroutine_threadsafe(notify(), asyncio_loop)
        # blocking the thread
        sleep(1)
        # make sure the task has finished
        task.result()

print("Main thread id {}".format(threading.get_ident()))

asyncio_loop = asyncio.get_event_loop()
executor = ThreadPoolExecutor(max_workers=1)
asyncio_loop.run_in_executor(executor, game_loop, asyncio_loop)

app = web.Application()

app.router.add_route('GET', '/connect', wshandler)
app.router.add_route('GET', '/', handle)

web.run_app(app)
<commit_msg>Update game_loop_thread.py<commit_after>import asyncio
from aiohttp import web

from concurrent.futures import ThreadPoolExecutor
import threading
from time import sleep

async def handle(request):
    index = open("index.html", 'rb')
    content = index.read()
    return web.Response(body=content, content_type='text/html')

tick = asyncio.Condition()

async def wshandler(request):
    ws = web.WebSocketResponse()
    await ws.prepare(request)

    recv_task = None
    tick_task = None
    while 1:
        if not recv_task:
            recv_task = asyncio.ensure_future(ws.receive())
        if not tick_task:
            await tick.acquire()
            tick_task = asyncio.ensure_future(tick.wait())

        done, pending = await asyncio.wait(
            [recv_task,
             tick_task],
            return_when=asyncio.FIRST_COMPLETED)

        if recv_task in done:
            msg = recv_task.result()
            if msg.tp == web.MsgType.text:
                print("Got message %s" % msg.data)
                ws.send_str("Pressed key code: {}".format(msg.data))
            elif msg.tp == web.MsgType.close or\
                 msg.tp == web.MsgType.error:
                break
            recv_task = None

        if tick_task in done:
            ws.send_str("game loop ticks")
            tick.release()
            tick_task = None

    return ws

def game_loop(asyncio_loop):
    print("Game loop thread id {}".format(threading.get_ident()))
    # a coroutine to run in main thread
    async def notify():
        print("Notify thread id {}".format(threading.get_ident()))
        await tick.acquire()
        tick.notify_all()
        tick.release()

    while 1:
        task = asyncio.run_coroutine_threadsafe(notify(), asyncio_loop)
        # blocking the thread
        sleep(1)
        # make sure the task has finished
        task.result()

print("Main thread id {}".format(threading.get_ident()))

asyncio_loop = asyncio.get_event_loop()
executor = ThreadPoolExecutor(max_workers=1)
asyncio_loop.run_in_executor(executor, game_loop, asyncio_loop)

app = web.Application()

app.router.add_route('GET', '/connect', wshandler)
app.router.add_route('GET', '/', handle)

web.run_app(app)
<|endoftext|>


<commit_before>from rest_framework import serializers

from . import models

class WidgetSerializer(serializers.ModelSerializer):
    class Meta:
        model = models.Widget
        fields = (
            "id",
            "owner",
            "name",
            "description",
            "public",
            "icon",
            "value",
            "timestamp",
            "type",
        )

    owner = serializers.ReadOnlyField(source="owner.username")
    type = serializers.ReadOnlyField(source="get_type_display")

class CommmentSerializer(serializers.ModelSerializer):
    class Meta:
        model = models.Comment
        fields = "__all__"

class SubscriptionSerializer(serializers.ModelSerializer):
    class Meta:
        model = models.Subscription
        fields = "__all__"

class SampleSerializer(serializers.ModelSerializer):
    class Meta:
        model = models.Sample
        fields = ("timestamp", "value")
<commit_msg>Fix title<commit_after>from rest_framework import serializers

from . import models

class WidgetSerializer(serializers.ModelSerializer):
    class Meta:
        model = models.Widget
        fields = (
            "id",
            "owner",
            "title",
            "description",
            "public",
            "icon",
            "value",
            "timestamp",
            "type",
        )

    owner = serializers.ReadOnlyField(source="owner.username")
    type = serializers.ReadOnlyField(source="get_type_display")

class CommmentSerializer(serializers.ModelSerializer):
    class Meta:
        model = models.Comment
        fields = "__all__"

class SubscriptionSerializer(serializers.ModelSerializer):
    class Meta:
        model = models.Subscription
        fields = "__all__"

class SampleSerializer(serializers.ModelSerializer):
    class Meta:
        model = models.Sample
        fields = ("timestamp", "value")
<|endoftext|>


<commit_before>#!/usr/bin/python

import json
import matplotlib.pyplot as plt
import sys, getopt

def main(argv):

    hybridfile = ""
    onlinefile = ""
    offlinefile = ""
    hybrid = False
    online = False
    offline = False
    hybrid_requests = None
    online_requests = None
    offline_requests = None

    try:
        opts, args = getopt.getopt(argv,"h",["hybrid_file=","online_file=","offline_file="])
    except getopt.GetoptError:
        print 'create_plots.py --hybrid_file <hybrid_inputfile> --online_file <online_inputfile> --offline_file <offline_inputfile>'
        sys.exit(2)
    for opt, arg in opts:
        if opt == '-h':
            print 'create_plots.py --hybrid_file <hybrid_inputfile> --online_file <online_inputfile> --offline_file <offline_inputfile>'
            sys.exit()
        elif opt in ("--hybrid_file"):
            hybridfile = arg
        elif opt in ("--online_file"):
            onlinefile = arg
        elif opt in ("--offline_file"):
            offlinefile = arg

    try:
        with open(hybridfile) as data_file:
            hybrid_requests = json.load(data_file)
            hybrid = True
    except:
        pass

    try:
        with open(onlinefile) as data_file:
            online_requests = json.load(data_file)
            online = True
    except:
        pass

    try:
        with open(offlinefile) as data_file:
            offline_requests = json.load(data_file)
            offline = True
    except:
        pass

    #Create mapped picture

    x=[]
    for i in xrange(0,len(hybrid_requests['mapped_requests'])):
        x.append(i)

    if hybrid:
        plt.plot(x, hybrid_requests['mapped_requests'])

    if online:
        plt.plot(x, online_requests['mapped_requests'])

    if offline:
        plt.plot(x, offline_requests['mapped_requests'])

    plt.title('Accepted incoming service requests')
    plt.ylabel('Accepted requests count')
    plt.xlabel('Incoming requests')
    plt.savefig("mapped_requests.png")
    plt.clf()

    #Create Running picture
    x=[]
    for i in xrange(0,len(hybrid_requests['running_requests'])):
        x.append(i)

    if hybrid:
        plt.plot(x, hybrid_requests['mapped_requests'])

    if online:
        plt.plot(x, online_requests['mapped_requests'])

    if offline:
        plt.plot(x, offline_requests['mapped_requests'])

    plt.plot(x,hybrid_requests['running_requests'])
    plt.title('Currently running (mapped) requests in the NFFG')
    plt.ylabel('Mapped requests count')
    plt.xlabel('Incoming requests')
    plt.savefig("running_requests.png")
    plt.clf()

    #Create refused picture
    x=[]
    for i in xrange(0,len(hybrid_requests['refused_requests'])):
        x.append(i)

    if hybrid:
        plt.plot(x, hybrid_requests['mapped_requests'])

    if online:
        plt.plot(x, online_requests['mapped_requests'])

    if offline:
        plt.plot(x, offline_requests['mapped_requests'])

    plt.plot(x,hybrid_requests['refused_requests'])
    plt.title('Refused requests during the simulation')
    plt.ylabel('Refused requests count')
    plt.xlabel('Incoming requests')
    plt.savefig("refused_requests.png")

if __name__ == "__main__":
   main(sys.argv[1:])<commit_msg>create plot bug fix<commit_after>#!/usr/bin/python

import json
import matplotlib.pyplot as plt
import sys, getopt

def main(argv):

    hybridfile = ""
    onlinefile = ""
    offlinefile = ""
    hybrid = False
    online = False
    offline = False
    hybrid_requests = None
    online_requests = None
    offline_requests = None

    try:
        opts, args = getopt.getopt(argv,"h",["hybrid_file=","online_file=","offline_file="])
    except getopt.GetoptError:
        print 'create_plots.py --hybrid_file <hybrid_inputfile> --online_file <online_inputfile> --offline_file <offline_inputfile>'
        sys.exit(2)
    for opt, arg in opts:
        if opt == '-h':
            print 'create_plots.py --hybrid_file <hybrid_inputfile> --online_file <online_inputfile> --offline_file <offline_inputfile>'
            sys.exit()
        elif opt in ("--hybrid_file"):
            hybridfile = arg
        elif opt in ("--online_file"):
            onlinefile = arg
        elif opt in ("--offline_file"):
            offlinefile = arg

    try:
        with open(hybridfile) as data_file:
            hybrid_requests = json.load(data_file)
            hybrid = True
    except:
        pass

    try:
        with open(onlinefile) as data_file:
            online_requests = json.load(data_file)
            online = True
    except:
        pass

    try:
        with open(offlinefile) as data_file:
            offline_requests = json.load(data_file)
            offline = True
    except:
        pass

    #Create mapped picture

    x=[]
    for i in xrange(0,len(hybrid_requests['running_requests'])):
        x.append(i)

    if hybrid:
        plt.plot(x, hybrid_requests['running_requests'])

    if online:
        plt.plot(x, online_requests['running_requests'])

    if offline:
        plt.plot(x, offline_requests['running_requests'])

    plt.title('Accepted incoming service requests')
    plt.ylabel('Accepted requests count')
    plt.xlabel('Incoming requests')
    plt.savefig("mapped_requests.png")
    plt.clf()

    #Create Running picture
    x=[]
    for i in xrange(0,len(hybrid_requests['running_requests'])):
        x.append(i)

    if hybrid:
        plt.plot(x, hybrid_requests['running_requests'])

    if online:
        plt.plot(x, online_requests['running_requests'])

    if offline:
        plt.plot(x, offline_requests['running_requests'])

    plt.plot(x,hybrid_requests['running_requests'])
    plt.title('Currently running (mapped) requests in the NFFG')
    plt.ylabel('Mapped requests count')
    plt.xlabel('Incoming requests')
    plt.savefig("running_requests.png")
    plt.clf()

    #Create refused picture
    x=[]
    for i in xrange(0,len(hybrid_requests['refused_requests'])):
        x.append(i)

    if hybrid:
        plt.plot(x, hybrid_requests['refused_requests'])

    if online:
        plt.plot(x, online_requests['refused_requests'])

    if offline:
        plt.plot(x, offline_requests['refused_requests'])

    plt.plot(x,hybrid_requests['refused_requests'])
    plt.title('Refused requests during the simulation')
    plt.ylabel('Refused requests count')
    plt.xlabel('Incoming requests')
    plt.savefig("refused_requests.png")

if __name__ == "__main__":
   main(sys.argv[1:])<|endoftext|>


<commit_before># vim: fdm=indent
# author:     Fabio Zanini
# date:       16/08/17
# content:    Dataset functions to cluster samples, features, and phenotypes.
# Modules
import numpy as np
import pandas as pd

from ..utils.cache import method_caches

# Classes / functions
class Cluster():
    '''Cluster samples, features, and phenotypes'''
    def __init__(self, dataset):
        '''Cluster samples, features, and phenotypes

        Args:
            dataset (Dataset): the dataset to analyze.
        '''
        self.dataset = dataset

    def kmeans(
            self,
            n_clusters,
            axis,
            phenotypes=(),
            random_state=0):
        '''K-Means clustering.

        Args:
            n_clusters (int): The number of clusters you want.
            axis (string): It must be 'samples' or 'features'. \
                    The Dataset.counts matrix is used and \
                    either samples or features are clustered.
            phenotypes (iterable of strings): Phenotypes to add to the \
                    features for joint clustering.
            log_features (bool): Whether to add pseudocounts and take a log \
                    of the feature counts before calculating distances.
            random_state (int): Set to the same int for deterministic results.

        Returns:
            pd.Series with the labels of the clusters.
        '''
        from sklearn.cluster import KMeans

        data = self.dataset.counts

        if phenotypes is not None:
            data = data.copy()
            for pheno in phenotypes:
                data.loc[pheno] = self.dataset.samplesheet.loc[:, pheno]

        if axis == 'samples':
            data = data.T
        elif axis == 'features':
            pass
        else:
            raise ValueError('axis must be "samples" or "features"')

        kmeans = (KMeans(n_clusters=n_clusters, random_state=random_state)
                  .fit(data.values))
        labels = pd.Series(kmeans.labels_, index=data.index, dtype='category')
        return labels

    def dbscan(
            self,
            axis,
            phenotypes=(),
            **kwargs):
        '''Density-Based Spatial Clustering of Applications with Noise.

        Args:
            axis (string): It must be 'samples' or 'features'. \
                    The Dataset.counts matrix is used and \
                    either samples or features are clustered.
            phenotypes (iterable of strings): Phenotypes to add to the \
                    features for joint clustering.
            log_features (bool): Whether to add pseudocounts and take a log \
                    of the feature counts before calculating distances.
            **kwargs: arguments passed to sklearn.cluster.DBSCAN.

        Returns:
            pd.Series with the labels of the clusters.
        '''
        from sklearn.cluster import DBSCAN

        data = self.dataset.counts

        if phenotypes is not None:
            data = data.copy()
            for pheno in phenotypes:
                data.loc[pheno] = self.dataset.samplesheet.loc[:, pheno]

        if axis == 'samples':
            data = data.T
        elif axis == 'features':
            pass
        else:
            raise ValueError('axis must be "samples" or "features"')

        kmeans = (DBSCAN(**kwargs)
                  .fit(data.values))
        labels = pd.Series(kmeans.labels_, index=data.index, dtype='category')
        return labels

    # NOTE: caching this one is tricky because it has non-kwargs AND it would
    # need a double cache, one for cells and one for features/phenotypes
    def hierarchical(
            self,
            axis,
            phenotypes=(),
            metric='correlation',
            method='average',
            log_features=True,
            optimal_ordering=False):
        '''Hierarchical clustering.

        Args:
            axis (string): It must be 'samples' or 'features'. \
                    The Dataset.counts matrix is used and \
                    either samples or features are clustered.
            phenotypes (iterable of strings): Phenotypes to add to the \
                    features for joint clustering.
            metric (string): Metric to calculate the distance matrix. Should \
                    be a string accepted by scipy.spatial.distance.pdist.
            method (string): Clustering method. Must be a string accepted by \
                    scipy.cluster.hierarchy.linkage.
            log_features (bool): Whether to add pseudocounts and take a log \
                    of the feature counts before calculating distances.
            optimal_ordering (bool): Whether to resort the linkage so that \
                    nearest neighbours have shortest distance. This may take \
                    longer than the clustering itself.
        Returns:
            dict with the linkage, distance matrix, and ordering.
        '''
        from scipy.spatial.distance import pdist
        from scipy.cluster.hierarchy import linkage, leaves_list

        if optimal_ordering:
            try:
                from polo import optimal_leaf_ordering
            except ImportError:
                raise ImportError(
                    'The package "polo" is needed for optimal leaf ordering')

        data = self.dataset.counts

        if log_features:
            data = np.log10(self.dataset.counts.pseudocount + data)

        if phenotypes is not None:
            data = data.copy()
            for pheno in phenotypes:
                data.loc[pheno] = self.dataset.samplesheet.loc[:, pheno]

        if axis == 'samples':
            data = data.T
        elif axis == 'features':
            pass
        else:
            raise ValueError('axis must be "samples" or "features"')

        Y = pdist(data.values, metric=metric)

        # Some metrics (e.g. correlation) give nan whenever the matrix has no
        # variation, default this to zero distance (e.g. two features that are
        # both total dropouts.
        Y = np.nan_to_num(Y)

        Z = linkage(Y, method=method)

        if optimal_ordering:
            Z = optimal_leaf_ordering(Z, Y)

        ids = data.index[leaves_list(Z)]

        return {
                'distance': Y,
                'linkage': Z,
                'leaves': ids,
                }
<commit_msg>API: do not log features in clustering by default<commit_after># vim: fdm=indent
# author:     Fabio Zanini
# date:       16/08/17
# content:    Dataset functions to cluster samples, features, and phenotypes.
# Modules
import numpy as np
import pandas as pd

from ..utils.cache import method_caches

# Classes / functions
class Cluster():
    '''Cluster samples, features, and phenotypes'''
    def __init__(self, dataset):
        '''Cluster samples, features, and phenotypes

        Args:
            dataset (Dataset): the dataset to analyze.
        '''
        self.dataset = dataset

    def kmeans(
            self,
            n_clusters,
            axis,
            phenotypes=(),
            random_state=0):
        '''K-Means clustering.

        Args:
            n_clusters (int): The number of clusters you want.
            axis (string): It must be 'samples' or 'features'. \
                    The Dataset.counts matrix is used and \
                    either samples or features are clustered.
            phenotypes (iterable of strings): Phenotypes to add to the \
                    features for joint clustering.
            log_features (bool): Whether to add pseudocounts and take a log \
                    of the feature counts before calculating distances.
            random_state (int): Set to the same int for deterministic results.

        Returns:
            pd.Series with the labels of the clusters.
        '''
        from sklearn.cluster import KMeans

        data = self.dataset.counts

        if phenotypes is not None:
            data = data.copy()
            for pheno in phenotypes:
                data.loc[pheno] = self.dataset.samplesheet.loc[:, pheno]

        if axis == 'samples':
            data = data.T
        elif axis == 'features':
            pass
        else:
            raise ValueError('axis must be "samples" or "features"')

        kmeans = (KMeans(n_clusters=n_clusters, random_state=random_state)
                  .fit(data.values))
        labels = pd.Series(kmeans.labels_, index=data.index, dtype='category')
        return labels

    def dbscan(
            self,
            axis,
            phenotypes=(),
            **kwargs):
        '''Density-Based Spatial Clustering of Applications with Noise.

        Args:
            axis (string): It must be 'samples' or 'features'. \
                    The Dataset.counts matrix is used and \
                    either samples or features are clustered.
            phenotypes (iterable of strings): Phenotypes to add to the \
                    features for joint clustering.
            log_features (bool): Whether to add pseudocounts and take a log \
                    of the feature counts before calculating distances.
            **kwargs: arguments passed to sklearn.cluster.DBSCAN.

        Returns:
            pd.Series with the labels of the clusters.
        '''
        from sklearn.cluster import DBSCAN

        data = self.dataset.counts

        if phenotypes is not None:
            data = data.copy()
            for pheno in phenotypes:
                data.loc[pheno] = self.dataset.samplesheet.loc[:, pheno]

        if axis == 'samples':
            data = data.T
        elif axis == 'features':
            pass
        else:
            raise ValueError('axis must be "samples" or "features"')

        kmeans = (DBSCAN(**kwargs)
                  .fit(data.values))
        labels = pd.Series(kmeans.labels_, index=data.index, dtype='category')
        return labels

    # NOTE: caching this one is tricky because it has non-kwargs AND it would
    # need a double cache, one for cells and one for features/phenotypes
    def hierarchical(
            self,
            axis,
            phenotypes=(),
            metric='correlation',
            method='average',
            log_features=False,
            optimal_ordering=False):
        '''Hierarchical clustering.

        Args:
            axis (string): It must be 'samples' or 'features'. \
                    The Dataset.counts matrix is used and \
                    either samples or features are clustered.
            phenotypes (iterable of strings): Phenotypes to add to the \
                    features for joint clustering.
            metric (string): Metric to calculate the distance matrix. Should \
                    be a string accepted by scipy.spatial.distance.pdist.
            method (string): Clustering method. Must be a string accepted by \
                    scipy.cluster.hierarchy.linkage.
            log_features (bool): Whether to add pseudocounts and take a log \
                    of the feature counts before calculating distances.
            optimal_ordering (bool): Whether to resort the linkage so that \
                    nearest neighbours have shortest distance. This may take \
                    longer than the clustering itself.
        Returns:
            dict with the linkage, distance matrix, and ordering.
        '''
        from scipy.spatial.distance import pdist
        from scipy.cluster.hierarchy import linkage, leaves_list

        if optimal_ordering:
            try:
                from polo import optimal_leaf_ordering
            except ImportError:
                raise ImportError(
                    'The package "polo" is needed for optimal leaf ordering')

        data = self.dataset.counts

        if log_features:
            data = np.log10(self.dataset.counts.pseudocount + data)

        if phenotypes is not None:
            data = data.copy()
            for pheno in phenotypes:
                data.loc[pheno] = self.dataset.samplesheet.loc[:, pheno]

        if axis == 'samples':
            data = data.T
        elif axis == 'features':
            pass
        else:
            raise ValueError('axis must be "samples" or "features"')

        Y = pdist(data.values, metric=metric)

        # Some metrics (e.g. correlation) give nan whenever the matrix has no
        # variation, default this to zero distance (e.g. two features that are
        # both total dropouts.
        Y = np.nan_to_num(Y)

        Z = linkage(Y, method=method)

        if optimal_ordering:
            Z = optimal_leaf_ordering(Z, Y)

        ids = data.index[leaves_list(Z)]

        return {
                'distance': Y,
                'linkage': Z,
                'leaves': ids,
                }
<|endoftext|>


<commit_before>"""

Revision ID: 0143_remove_reply_to
Revises: 0142_validate_constraint
Create Date: 2017-11-20 14:35:16.826398

"""
from alembic import op
import sqlalchemy as sa

revision = '0143_remove_reply_to'
down_revision = '0142_validate_constraint'

def upgrade():
    op.drop_column('services', 'letter_contact_block')
    op.drop_column('services', 'reply_to_email_address')
    op.drop_column('services_history', 'letter_contact_block')
    op.drop_column('services_history', 'reply_to_email_address')

def downgrade():
    op.add_column('services_history', sa.Column('reply_to_email_address', sa.TEXT(), autoincrement=False, nullable=True))
    op.add_column('services_history', sa.Column('letter_contact_block', sa.TEXT(), autoincrement=False, nullable=True))
    op.add_column('services', sa.Column('reply_to_email_address', sa.TEXT(), autoincrement=False, nullable=True))
    op.add_column('services', sa.Column('letter_contact_block', sa.TEXT(), autoincrement=False, nullable=True))
<commit_msg>Need to do the db migration in the next deploy<commit_after><|endoftext|>


<commit_before># -*- coding: utf-8 -*-
#
#
#

from src.probabilistic_seq_match import sequences_match_prob
from src.probabilistic_seq_match import base_prob
import sys
import os
from shutil import rmtree
import glob
import gzip
from Bio import SeqIO
from Bio.Seq import Seq
from Bio.SeqRecord import SeqRecord
from Bio.Alphabet import SingleLetterAlphabet
from operator import itemgetter
import concurrent.futures as cf

def run(args):

    # Precompute base probabilities for phredscores up to 50
    base_prob_precompute = {}
    for score in range(1, 51):
        base_prob_precompute[score] = base_prob(score)

    # Check that the multiplexed path exists
    multiplexed_dir = os.path.join(args.inDir, "multiplexed")
    if not os.path.exists(multiplexed_dir):
        sys.exit("Directory '<inDir>/multiplexed' does not exist. Re-run with"
                 " different <inDir>")

    # Create out directory
    out_dir = "demultiplexed"
    if args.uniqID != None:
        out_dir += "_{0}".format(args.uniqID)
    out_dir = os.path.join(args.inDir, out_dir)
    create_folder(out_dir)

    # Initiate multiplexed class
    multiplexed = Multiplex(multiplexed_dir)

    # Initiate sample sheet and read possible indexes
    sampleSheet = SampleSheet(args.sampleSheet)
    sampleSheet.parse(args.indexQual)

    # Check that there are the same number of indexes in sample sheet and
    # multiplexed fastqs
    if sampleSheet.is_dualindexed != multiplexed.is_dualindexed:
        sys.exit("Error: Different number of indexes in sampleSheet and "
                 "multiplexed reads. Exiting!")

    # Open output class for each sample, and a not_assigned group
    sample_out = {}
    for sample in list(sampleSheet.sample_indexes.keys()) + ['not_assigned']:
        sample_out[sample] = Sample(sample, out_dir, multiplexed.is_pairend,
                                    multiplexed.is_dualindexed)

    # Send each read/barcode record to futures to match up to sample
    with cf.ProcessPoolExecutor(max_workers=args.numCPU) as executor:
        c = 1
        # Map read/barcode records
        for output in executor.map(futures_barcode_to_indexes,
            futures_iterate_reads(multiplexed, sampleSheet,
                base_prob_precompute, args.minProb)):
            # Unpack output
            ((read_records, barcode_records), sample, prob, _) = output
            # Write record to correct sample file
            sample_out[sample].write(read_records, barcode_records)

            # Update progress
            if c % 1000 == 0:
                print(c)
            c += 1

    return 0

def futures_iterate_reads(multiplexed, sampleSheet, base_prob_precompute,
                          min_prob):
    """ Returns an iterator that contains everything needed for futures.
    """
    for combined_record in multiplexed.iterate():
        yield (combined_record, sampleSheet, base_prob_precompute, min_prob)

def futures_barcode_to_indexes(variables):
    """ Compares the reads barcodes to sample indexes and returns matching
        sample name.
    """
    # Unpack variables
    (combined_record, sampleSheet,
     base_prob_precompute, min_prob) = variables
    # Get barcode records
    _, barcode_records = combined_record
    # Find sample
    b1_header, sample, prob = match_barcode_to_indexes(barcode_records,
        sampleSheet, base_prob_precompute, min_prob)
    if sample == None:
        sample = 'not_assigned'
    # Append probability to barcode1 header
    b1_header = "{0}_{1}".format(b1_header, prob)
    # Change header
    combined_record[1][0].id = b1_header

    return combined_record, sample, prob, b1_header

def match_barcode_to_indexes(barcode_records, sampleSheet, base_prob_precompute,
                             min_prob):
    """ For the barcode pair, caluclates probability of a match against each set
        of indexes
    """

    index_probs = {}

    for sample_name in sampleSheet.sample_indexes:

        index_records = sampleSheet.sample_indexes[sample_name]

        # Calculate the match probability for barcode 1
        b1_prob = sequences_match_prob(index_records[0].seq,
            index_records[0].letter_annotations['phred_quality'],
            barcode_records[0].seq,
            barcode_records[0].letter_annotations['phred_quality'],
            base_prob_precompute, 0)

        # Do for second barcode if present
        if sampleSheet.is_dualindexed:
            # Skip if already below the threshold, else assign same prob as b1
            if b1_prob >= min_prob:
                b2_prob = sequences_match_prob(index_records[1].seq,
                           index_records[1].letter_annotations['phred_quality'],
                           barcode_records[1].seq,
                           barcode_records[1].letter_annotations['phred_quality'],
                           base_prob_precompute, 0)
            else:
                b2_prob = b1_prob

        # Caluclate combined probability
        if sampleSheet.is_dualindexed:
            overall_prob = b1_prob * b2_prob
        else:
            overall_prob = b1_prob

        index_probs[sample_name] = overall_prob

    sorted_probs = sorted(index_probs.items(), key=itemgetter(1),
                          reverse=True)

    # Return header, sample, prob
    header = barcode_records[0].id
    if sorted_probs[0][1] > min_prob:
        return header, sorted_probs[0][0], sorted_probs[0][1]
    else:
        return header, None, sorted_probs[0][1]

class Sample:
    # Class for each possible sample. 1) Holds the output directory for that
    # sample. 2) Opens handles. 3) Writes record to sample.

    def __init__(self, name, out_dir, is_pe, id_dual):
        self.read_paths = []
        self.barcode_paths = []
        self.read_handles = None
        self.barcode_handles = None

        # Create directory for sample
        name = name.replace(' ', '_')
        self.sample_dir = os.path.join(out_dir, name)
        create_folder(self.sample_dir)

        # Create read paths
        self.read_paths.append(os.path.join(self.sample_dir,
            '{0}.R1.fastq.gz'.format(name)))
        if is_pe:
            self.read_paths.append(os.path.join(self.sample_dir,
                '{0}.R2.fastq.gz'.format(name)))
        # Create barcode paths
        self.barcode_paths.append(os.path.join(self.sample_dir,
            '{0}.barcode_1.fastq.gz'.format(name)))
        if id_dual:
            self.barcode_paths.append(os.path.join(self.sample_dir,
                '{0}.barcode_2.fastq.gz'.format(name)))

    def open_handles(self):
        """ For the reads and barcodes, opens output handles.
        """
        self.read_handles = [get_handle(read_path, 'wt') for read_path
                             in self.read_paths]
        self.barcode_handles = [get_handle(barcode_path, 'wt') for barcode_path
                                in self.barcode_paths]
        return 0

    def write(self, read_records, barcode_records):
        """ Writes the demultiplexed read and barcode records to sample file.
        """
        # Open handles if not open
        if self.read_handles == None:
            self.open_handles()
        # Write read records
        for i in range(len(read_records)):
            SeqIO.write(read_records[i], self.read_handles[i], "fastq")
        # Write barcode records
        for i in range(len(barcode_records)):
            SeqIO.write(barcode_records[i], self.barcode_handles[i], "fastq")
        return 0

    def close_handles():
        """ Closes any open handles.
        """
        if self.read_handles != None:
            for handle in self.read_handles + self.barcode_handles:
                handle.close()
        return 0

class SampleSheet:
    # Class to hold the sample sheet and retrieve indexes from it.

    def __init__(self, path):
        self.path = path

    def parse(self, index_qual):
        """ Parses the sample sheet to retrieve the indexes for each sample.
        """
        sample_indexes = {}
        with open(self.path, 'r') as in_h:
            # Skip to line after [Data]
            line = in_h.readline()
            while not line.startswith('[Data]'):
                line = in_h.readline()
            # Get header
            header = in_h.readline().rstrip().lower().split(',')
            col_ind = dict(zip(header, range(len(header))))
            # Save whether it is dual indexed
            if "index2" in col_ind.keys():
                self.is_dualindexed = True
            else:
                self.is_dualindexed = False
            # Get indexes
            for line in in_h:
                # Break if EOF
                if line.strip() == "":
                    break
                # Get info
                parts = line.rstrip().split(',')
                sample_name = parts[col_ind['sample_name']]
                # Get first index
                index1 = parts[col_ind['index']]
                sample_indexes[sample_name] = [index1]
                # Get second index
                if self.is_dualindexed:
                    index2 = parts[col_ind['index2']]
                sample_indexes[sample_name].append(index2)

        # Convert indexes to seqIO seqRecords
        self.sample_indexes = self.convert_index_to_seqRecord(sample_indexes,
            index_qual)

        return 0

    def convert_index_to_seqRecord(self, sample_indexes, index_qual):
        """ Converts each index sequence to a seqIO seqRecord.
        """
        # For each sample
        for sample in sample_indexes:
            # For each index
            for i in range(len(sample_indexes[sample])):
                raw_seq = sample_indexes[sample][i]
                # Convert to seqRecord
                sample_indexes[sample][i] = SeqRecord(
                    Seq(raw_seq, SingleLetterAlphabet()))
                # Add quality information
                sample_indexes[sample][i].letter_annotations['phred_quality'] \
                    = [index_qual] * len(raw_seq)
        return sample_indexes

class Multiplex:
    # Class for the folder of multiplexed reads + barcodes

    def __init__(self, folder):
        """ Make list of read and barcode files.
        """
        self.dir = folder
        # Get list of read and barcode paths
        self.read_paths = []
        self.barcode_paths = []
        for fastq in sorted(glob.glob(os.path.join(folder, "*.fastq*"))):
            if "barcode_" in os.path.split(fastq)[1]:
                self.barcode_paths.append(fastq)
            else:
                self.read_paths.append(fastq)
        # Save whether pairend
        if len(self.read_paths) == 1:
            self.is_pairend = False
        elif len(self.read_paths) == 2:
            self.is_pairend = True
        else:
            sys.exit("There must be 1 or 2 input read fastqs, not {0}".format(
                len(self.read_paths)))
        # Save whether dualindex
        if len(self.barcode_paths) == 1:
            self.is_dualindexed = False
        elif len(self.barcode_paths) == 2:
            self.is_dualindexed = True
        else:
            sys.exit("There must be 1 or 2 input barcode fastqs, not"
                     " {0}".format(len(self.barcode_paths)))
        return None

    def open_handles(self):
        """ Opens the file names for reading.
        """
        read_handles = [get_handle(filen, 'r') for filen in self.read_paths]
        barcode_handles = [get_handle(filen, 'r') for filen
                           in self.barcode_paths]
        return read_handles, barcode_handles

    def open_seqIO_iterators(self, read_handles, barcode_handles):
        """ Opens fastq iterators using biopythons SeqIO
        """
        # Open iterators for each handle
        read_iterators = [SeqIO.parse(handle, "fastq") for handle
                          in read_handles]
        barcode_iterators = [SeqIO.parse(handle, "fastq") for handle
                             in barcode_handles]
        return read_iterators, barcode_iterators

    def iterate(self):
        """ Loads the reads and barcode fastqs and yields 1 set at a time.
        """
        # Open handles
        read_handles, barcode_handles = self.open_handles()
        # Open iterators for each handle
        read_iterators, barcode_iterators = self.open_seqIO_iterators(
            read_handles, barcode_handles)

        # Iterate through records
        for r1_record in read_iterators[0]:

            # Get read records
            read_records = [r1_record]
            if self.is_pairend:
                read_records.append(next(read_iterators[1]))
            # Get barcode records
            barcode_records = [next(barcode_iterators[0])]
            if self.is_dualindexed:
                barcode_records.append(next(barcode_iterators[1]))

            # Check that they all have the same title
            titles = [record.id for record in read_records + barcode_records]
            if len(set(titles)) > 1:
                sys.exit('Reads and/or barcodes are not in sync\n'
                         '{0}'.format(titles))

            yield [read_records, barcode_records]

        # Close handles
        for handle in read_handles + barcode_handles:
            handle.close()

def create_folder(folder):
    """ Check out folder exists and create a new one.
    """
    # Check if it exists
    if os.path.exists(folder):
        response = input('{0} exists. Would you like to overwrite it? [y/n] '.format(folder))
        if response == 'y':
            rmtree(folder)
        else:
            sys.exit()
    os.makedirs(folder)
    return folder

def get_handle(filen, rw):
    """ Returns file handle using gzip if file ends in .gz
    """
    if filen.split('.')[-1] == 'gz':
        return gzip.open(filen, rw)
    else:
        return open(filen, rw)
<commit_msg>switch to non futures<commit_after># -*- coding: utf-8 -*-
#
#
#

from src.probabilistic_seq_match import sequences_match_prob
from src.probabilistic_seq_match import base_prob
import sys
import os
from shutil import rmtree
import glob
import gzip
from Bio import SeqIO
from Bio.Seq import Seq
from Bio.SeqRecord import SeqRecord
from Bio.Alphabet import SingleLetterAlphabet
from operator import itemgetter
import concurrent.futures as cf

def run(args):

    # Precompute base probabilities for phredscores up to 50
    base_prob_precompute = {}
    for score in range(1, 51):
        base_prob_precompute[score] = base_prob(score)

    # Check that the multiplexed path exists
    multiplexed_dir = os.path.join(args.inDir, "multiplexed")
    if not os.path.exists(multiplexed_dir):
        sys.exit("Directory '<inDir>/multiplexed' does not exist. Re-run with"
                 " different <inDir>")

    # Create out directory
    out_dir = "demultiplexed"
    if args.uniqID != None:
        out_dir += "_{0}".format(args.uniqID)
    out_dir = os.path.join(args.inDir, out_dir)
    create_folder(out_dir)

    # Initiate multiplexed class
    multiplexed = Multiplex(multiplexed_dir)

    # Initiate sample sheet and read possible indexes
    sampleSheet = SampleSheet(args.sampleSheet)
    sampleSheet.parse(args.indexQual)

    # Check that there are the same number of indexes in sample sheet and
    # multiplexed fastqs
    if sampleSheet.is_dualindexed != multiplexed.is_dualindexed:
        sys.exit("Error: Different number of indexes in sampleSheet and "
                 "multiplexed reads. Exiting!")

    # Open output class for each sample, and a not_assigned group
    sample_out = {}
    for sample in list(sampleSheet.sample_indexes.keys()) + ['not_assigned']:
        sample_out[sample] = Sample(sample, out_dir, multiplexed.is_pairend,
                                    multiplexed.is_dualindexed)

    # Version that doesnt use concurrent.futures
    c = 1
    for variables in futures_iterate_reads(multiplexed, sampleSheet,
                                           base_prob_precompute, args.minProb):
        output = futures_barcode_to_indexes(variables)
        # Unpack output
        ((read_records, barcode_records), sample, prob, _) = output
        # Write record to correct sample file
        sample_out[sample].write(read_records, barcode_records)
        # Update progress
        if c % 1000 == 0:
            print(c)
        c += 1

    """
    # Send each read/barcode record to futures to match up to sample
    with cf.ProcessPoolExecutor(max_workers=args.numCPU) as executor:
        c = 1
        # Map read/barcode records
        for output in executor.map(futures_barcode_to_indexes,
            futures_iterate_reads(multiplexed, sampleSheet,
                base_prob_precompute, args.minProb)):
            # Unpack output
            ((read_records, barcode_records), sample, prob, _) = output
            # Write record to correct sample file
            sample_out[sample].write(read_records, barcode_records)

            # Update progress
            if c % 1000 == 0:
                print(c)
            c += 1
    """

    return 0

def futures_iterate_reads(multiplexed, sampleSheet, base_prob_precompute,
                          min_prob):
    """ Returns an iterator that contains everything needed for futures.
    """
    for combined_record in multiplexed.iterate():
        yield (combined_record, sampleSheet, base_prob_precompute, min_prob)

def futures_barcode_to_indexes(variables):
    """ Compares the reads barcodes to sample indexes and returns matching
        sample name.
    """
    # Unpack variables
    (combined_record, sampleSheet,
     base_prob_precompute, min_prob) = variables
    # Get barcode records
    _, barcode_records = combined_record
    # Find sample
    b1_header, sample, prob = match_barcode_to_indexes(barcode_records,
        sampleSheet, base_prob_precompute, min_prob)
    if sample == None:
        sample = 'not_assigned'
    # Append probability to barcode1 header
    b1_header = "{0}_{1}".format(b1_header, prob)
    # Change header
    combined_record[1][0].id = b1_header

    return combined_record, sample, prob, b1_header

def match_barcode_to_indexes(barcode_records, sampleSheet, base_prob_precompute,
                             min_prob):
    """ For the barcode pair, caluclates probability of a match against each set
        of indexes
    """

    index_probs = {}

    for sample_name in sampleSheet.sample_indexes:

        index_records = sampleSheet.sample_indexes[sample_name]

        # Calculate the match probability for barcode 1
        b1_prob = sequences_match_prob(index_records[0].seq,
            index_records[0].letter_annotations['phred_quality'],
            barcode_records[0].seq,
            barcode_records[0].letter_annotations['phred_quality'],
            base_prob_precompute, 0)

        # Do for second barcode if present
        if sampleSheet.is_dualindexed:
            # Skip if already below the threshold, else assign same prob as b1
            if b1_prob >= min_prob:
                b2_prob = sequences_match_prob(index_records[1].seq,
                           index_records[1].letter_annotations['phred_quality'],
                           barcode_records[1].seq,
                           barcode_records[1].letter_annotations['phred_quality'],
                           base_prob_precompute, 0)
            else:
                b2_prob = b1_prob

        # Caluclate combined probability
        if sampleSheet.is_dualindexed:
            overall_prob = b1_prob * b2_prob
        else:
            overall_prob = b1_prob

        index_probs[sample_name] = overall_prob

    sorted_probs = sorted(index_probs.items(), key=itemgetter(1),
                          reverse=True)

    # Return header, sample, prob
    header = barcode_records[0].id
    if sorted_probs[0][1] > min_prob:
        return header, sorted_probs[0][0], sorted_probs[0][1]
    else:
        return header, None, sorted_probs[0][1]

class Sample:
    # Class for each possible sample. 1) Holds the output directory for that
    # sample. 2) Opens handles. 3) Writes record to sample.

    def __init__(self, name, out_dir, is_pe, id_dual):
        self.read_paths = []
        self.barcode_paths = []
        self.read_handles = None
        self.barcode_handles = None

        # Create directory for sample
        name = name.replace(' ', '_')
        self.sample_dir = os.path.join(out_dir, name)
        create_folder(self.sample_dir)

        # Create read paths
        self.read_paths.append(os.path.join(self.sample_dir,
            '{0}.R1.fastq.gz'.format(name)))
        if is_pe:
            self.read_paths.append(os.path.join(self.sample_dir,
                '{0}.R2.fastq.gz'.format(name)))
        # Create barcode paths
        self.barcode_paths.append(os.path.join(self.sample_dir,
            '{0}.barcode_1.fastq.gz'.format(name)))
        if id_dual:
            self.barcode_paths.append(os.path.join(self.sample_dir,
                '{0}.barcode_2.fastq.gz'.format(name)))

    def open_handles(self):
        """ For the reads and barcodes, opens output handles.
        """
        self.read_handles = [get_handle(read_path, 'wt') for read_path
                             in self.read_paths]
        self.barcode_handles = [get_handle(barcode_path, 'wt') for barcode_path
                                in self.barcode_paths]
        return 0

    def write(self, read_records, barcode_records):
        """ Writes the demultiplexed read and barcode records to sample file.
        """
        # Open handles if not open
        if self.read_handles == None:
            self.open_handles()
        # Write read records
        for i in range(len(read_records)):
            SeqIO.write(read_records[i], self.read_handles[i], "fastq")
        # Write barcode records
        for i in range(len(barcode_records)):
            SeqIO.write(barcode_records[i], self.barcode_handles[i], "fastq")
        return 0

    def close_handles():
        """ Closes any open handles.
        """
        if self.read_handles != None:
            for handle in self.read_handles + self.barcode_handles:
                handle.close()
        return 0

class SampleSheet:
    # Class to hold the sample sheet and retrieve indexes from it.

    def __init__(self, path):
        self.path = path

    def parse(self, index_qual):
        """ Parses the sample sheet to retrieve the indexes for each sample.
        """
        sample_indexes = {}
        with open(self.path, 'r') as in_h:
            # Skip to line after [Data]
            line = in_h.readline()
            while not line.startswith('[Data]'):
                line = in_h.readline()
            # Get header
            header = in_h.readline().rstrip().lower().split(',')
            col_ind = dict(zip(header, range(len(header))))
            # Save whether it is dual indexed
            if "index2" in col_ind.keys():
                self.is_dualindexed = True
            else:
                self.is_dualindexed = False
            # Get indexes
            for line in in_h:
                # Break if EOF
                if line.strip() == "":
                    break
                # Get info
                parts = line.rstrip().split(',')
                sample_name = parts[col_ind['sample_name']]
                # Get first index
                index1 = parts[col_ind['index']]
                sample_indexes[sample_name] = [index1]
                # Get second index
                if self.is_dualindexed:
                    index2 = parts[col_ind['index2']]
                sample_indexes[sample_name].append(index2)

        # Convert indexes to seqIO seqRecords
        self.sample_indexes = self.convert_index_to_seqRecord(sample_indexes,
            index_qual)

        return 0

    def convert_index_to_seqRecord(self, sample_indexes, index_qual):
        """ Converts each index sequence to a seqIO seqRecord.
        """
        # For each sample
        for sample in sample_indexes:
            # For each index
            for i in range(len(sample_indexes[sample])):
                raw_seq = sample_indexes[sample][i]
                # Convert to seqRecord
                sample_indexes[sample][i] = SeqRecord(
                    Seq(raw_seq, SingleLetterAlphabet()))
                # Add quality information
                sample_indexes[sample][i].letter_annotations['phred_quality'] \
                    = [index_qual] * len(raw_seq)
        return sample_indexes

class Multiplex:
    # Class for the folder of multiplexed reads + barcodes

    def __init__(self, folder):
        """ Make list of read and barcode files.
        """
        self.dir = folder
        # Get list of read and barcode paths
        self.read_paths = []
        self.barcode_paths = []
        for fastq in sorted(glob.glob(os.path.join(folder, "*.fastq*"))):
            if "barcode_" in os.path.split(fastq)[1]:
                self.barcode_paths.append(fastq)
            else:
                self.read_paths.append(fastq)
        # Save whether pairend
        if len(self.read_paths) == 1:
            self.is_pairend = False
        elif len(self.read_paths) == 2:
            self.is_pairend = True
        else:
            sys.exit("There must be 1 or 2 input read fastqs, not {0}".format(
                len(self.read_paths)))
        # Save whether dualindex
        if len(self.barcode_paths) == 1:
            self.is_dualindexed = False
        elif len(self.barcode_paths) == 2:
            self.is_dualindexed = True
        else:
            sys.exit("There must be 1 or 2 input barcode fastqs, not"
                     " {0}".format(len(self.barcode_paths)))
        return None

    def open_handles(self):
        """ Opens the file names for reading.
        """
        read_handles = [get_handle(filen, 'r') for filen in self.read_paths]
        barcode_handles = [get_handle(filen, 'r') for filen
                           in self.barcode_paths]
        return read_handles, barcode_handles

    def open_seqIO_iterators(self, read_handles, barcode_handles):
        """ Opens fastq iterators using biopythons SeqIO
        """
        # Open iterators for each handle
        read_iterators = [SeqIO.parse(handle, "fastq") for handle
                          in read_handles]
        barcode_iterators = [SeqIO.parse(handle, "fastq") for handle
                             in barcode_handles]
        return read_iterators, barcode_iterators

    def iterate(self):
        """ Loads the reads and barcode fastqs and yields 1 set at a time.
        """
        # Open handles
        read_handles, barcode_handles = self.open_handles()
        # Open iterators for each handle
        read_iterators, barcode_iterators = self.open_seqIO_iterators(
            read_handles, barcode_handles)

        # Iterate through records
        for r1_record in read_iterators[0]:

            # Get read records
            read_records = [r1_record]
            if self.is_pairend:
                read_records.append(next(read_iterators[1]))
            # Get barcode records
            barcode_records = [next(barcode_iterators[0])]
            if self.is_dualindexed:
                barcode_records.append(next(barcode_iterators[1]))

            # Check that they all have the same title
            titles = [record.id for record in read_records + barcode_records]
            if len(set(titles)) > 1:
                sys.exit('Reads and/or barcodes are not in sync\n'
                         '{0}'.format(titles))

            yield [read_records, barcode_records]

        # Close handles
        for handle in read_handles + barcode_handles:
            handle.close()

def create_folder(folder):
    """ Check out folder exists and create a new one.
    """
    # Check if it exists
    if os.path.exists(folder):
        response = input('{0} exists. Would you like to overwrite it? [y/n] '.format(folder))
        if response == 'y':
            rmtree(folder)
        else:
            sys.exit()
    os.makedirs(folder)
    return folder

def get_handle(filen, rw):
    """ Returns file handle using gzip if file ends in .gz
    """
    if filen.split('.')[-1] == 'gz':
        return gzip.open(filen, rw)
    else:
        return open(filen, rw)
<|endoftext|>