<commit_before># -*- coding: utf-8 -*-
"""Database module, including the SQLAlchemy database object and DB-related utilities."""
from pprint import pformat

from sqlalchemy import desc, or_
from sqlalchemy.sql.sqltypes import Date, DateTime
from werkzeug import cached_property
from flask import current_app
from walle.service.extensions import db
from walle.service.utils import basestring
from walle.service.utils import datetime_str_to_obj, date_str_to_obj

# Alias common SQLAlchemy names
Column = db.Column
relationship = db.relationship

OPERATOR_FUNC_DICT = {
    '=': (lambda cls, k, v: getattr(cls, k) == v),
    '==': (lambda cls, k, v: getattr(cls, k) == v),
    'eq': (lambda cls, k, v: getattr(cls, k) == v),
    '!=': (lambda cls, k, v: getattr(cls, k) != v),
    'ne': (lambda cls, k, v: getattr(cls, k) != v),
    'neq': (lambda cls, k, v: getattr(cls, k) != v),
    '>': (lambda cls, k, v: getattr(cls, k) > v),
    'gt': (lambda cls, k, v: getattr(cls, k) > v),
    '>=': (lambda cls, k, v: getattr(cls, k) >= v),
    'gte': (lambda cls, k, v: getattr(cls, k) >= v),
    '<': (lambda cls, k, v: getattr(cls, k) < v),
    'lt': (lambda cls, k, v: getattr(cls, k) < v),
    '<=': (lambda cls, k, v: getattr(cls, k) <= v),
    'lte': (lambda cls, k, v: getattr(cls, k) <= v),
    'or': (lambda cls, k, v: or_(getattr(cls, k) == value for value in v)),
    'in': (lambda cls, k, v: getattr(cls, k).in_(v)),
    'nin': (lambda cls, k, v: ~getattr(cls, k).in_(v)),
    'like': (lambda cls, k, v: getattr(cls, k).like('%%%s%%' % (v))),
    'nlike': (lambda cls, k, v: ~getattr(cls, k).like(v)),
    '+': (lambda cls, k, v: getattr(cls, k) + v),
    'incr': (lambda cls, k, v: getattr(cls, k) + v),
    '-': (lambda cls, k, v: getattr(cls, k) - v),
    'decr': (lambda cls, k, v: getattr(cls, k) - v),
}

def parse_operator(cls, filter_name_dict):
    """ 用来返回sqlalchemy query对象filter使用的表达式
    Args:
        filter_name_dict (dict): 过滤条件dict
        {
            'last_name': {'eq': 'wang'},    # 如果是dic使用key作为操作符
            'age': {'>': 12}
        }
    Returns:
        binary_expression_list (lambda list)
    """

    def _change_type(cls, field, value):
        """ 有些表字段比如DateTime类型比较的时候需要转换类型，
        前端传过来的都是字符串，Date等类型没法直接相比较，需要转成Date类型
        Args:
            cls (class): Model class
            field (str): Model class field
            value (str): value need to compare
        """
        field_type = getattr(cls, field).type
        if isinstance(field_type, Date):
            return date_str_to_obj(value)
        elif isinstance(field_type, DateTime):
            return datetime_str_to_obj(value)
        else:
            return value

    binary_expression_list = []
    for field, op_dict in list(filter_name_dict.items()):
        for op, op_val in list(op_dict.items()):
            op_val = _change_type(cls, field, op_val)
            if op in OPERATOR_FUNC_DICT:
                binary_expression_list.append(
                        OPERATOR_FUNC_DICT[op](cls, field, op_val)
                )
    return binary_expression_list

class CRUDMixin(object):
    """Mixin that adds convenience methods for
    CRUD (create, read, update, delete) operations."""

    @classmethod
    def create(cls, **kwargs):
        """Create a new record and save it the database."""
        instance = cls(**kwargs)
        return instance.save()

    @classmethod
    def create_from_dict(cls, d):
        """Create a new record and save it the database."""
        assert isinstance(d, dict)
        instance = cls(**d)
        return instance.save()

    def update(self, commit=True, **kwargs):
        """Update specific fields of a record."""
        for attr, value in list(kwargs.items()):
            setattr(self, attr, value)
        return commit and self.save() or self

    def save(self, commit=True):
        """Save the record."""
        db.session.add(self)
        if commit:
            try:
                db.session.commit()
            except Exception as e:
                current_app.logger.info(e)
                db.session.rollback()
        return self

    def delete(self, commit=True):
        """Remove the record from the database."""
        db.session.delete(self)
        if commit:
            try:
                db.session.commit()
            except:
                db.session.rollback()
        return self

    def to_dict(self, fields_list=None):
        """
        Args:
            fields (str list): 指定返回的字段
        """
        column_list = fields_list or [
            column.name for column in self.__table__.columns
            ]
        return {
            column_name: getattr(self, column_name)
            for column_name in column_list
            }

    @classmethod
    def create_or_update(cls, query_dict, update_dict=None):
        instance = db.session.query(cls).filter_by(**query_dict).first()
        if instance:  # update
            if update_dict is not None:
                return instance.update(**update_dict)
            else:
                return instance
        else:  # create new instance
            query_dict.update(update_dict or {})
            return cls.create(**query_dict)

    @classmethod
    def query_paginate(cls, page=1, limit=20, fields=None, order_by_list=[('id', 'desc')],
                       filter_name_dict=None):
        """ 通用的分页查询函数
        Args:
            page (int): 页数
            limit (int): 每页个数
            order_by_list (tuple list): 用来指定排序的字段，可以传多个
                [ ('id', 1), ('name', -1) ]，1表示正序，-1表示逆序
                or
                [ ('id', 'asc'), ('name', 'desc') ]，1表示正序，-1表示逆序

            filter_name_dict (dict): 过滤条件，使用字典表示，使用字段名作为key，value
                是{'operator': to_compare_value}, e.g.:
                {
                    'last_name': {'eq': 'wang'},  # 如果是dic使用key作为操作符
                    'age': {'>': 12}
                }

        Returns:
            if fields is not None:
                (keytuple_list, total_cnt) (tuple)
            else:
                (instance_list, total_cnt) (tuple)

        前段查询参数规范：
        request.args 示例：
        ImmutableMultiDict([('limit', '10'), ('page', '1'), ('filter', '[{"field":"name","op":"eq","q":"g"},{
        "field":"id","op":"gt","q":"5"
        }]')])

        page: 页码
        limit: 每页限制
        order: 顺序，取值"asc" "desc". """'name', 'asc', 'model', 'desc'"""
        fields: 需要返回的字段
        filter: 过滤条件：
        {
            field: 需要过滤的字段
            op: 过滤操作符,支持"eq","neq","gt","gte","lt","lte","in","nin","like"
            q: 过滤值
        }
        """
        fields = (
            [getattr(cls, column) for column in fields] if fields is not None
            else None
        )
        if fields:
            query = db.session.query(*fields)
        else:
            query = db.session.query(cls)
        if order_by_list:
            for (field, order) in order_by_list:
                query = (
                    query.order_by(getattr(cls, field)) if order == 1 else
                    query.order_by(desc(getattr(cls, field)))
                )

        if filter_name_dict:
            p = parse_operator(cls, filter_name_dict)
            query = query.filter(*p)

        total_cnt = query.count()
        start = (page - 1) * limit
        query = query.offset(start).limit(limit)
        instance_or_keytuple_list = query.all()
        return instance_or_keytuple_list, total_cnt

    @classmethod
    def dump_schema(cls, items, fields, schema_class):
        """ 用来序列化从数据库查询出来的对象
        Args:
            items (instance list): obj list query from mysql
            fields (str list): fields need to dump
            schema_class (marshmallow.Schema): marshmallow.Schema class
        Returns:
            items, err
        """
        fields = (
            fields if fields else list(schema_class._declared_fields.keys())
        )
        schema = schema_class(many=True, only=fields)
        items, err = schema.dump(items)
        return items, err

    @classmethod
    def query_paginate_and_dump_schema(cls, page=1, limit=20, fields=None,
                                       order_by_list=None,
                                       filter_name_dict=None,
                                       schema_class=None):
        """ 分页查询并且返回dump后的对象，可以解决大部分查询问题 """
        assert schema_class
        items, total_cnt = cls.query_paginate(
                page, limit, fields, order_by_list, filter_name_dict
        )
        items, err = cls.dump_schema(items, fields, schema_class)
        return items, total_cnt

    def __repr__(self):
        return pformat(self.to_dict())

    @cached_property
    def column_name_set(self):
        return set([column.name for column in self.__table__.columns])

    @classmethod
    def get_common_fields(cls, fields=None):
        """ 防止传过来的fields含有该Model没有的字段 """
        if not fields:
            return []
        table_fields_set = set(
                [column.name for column in cls.__table__.columns]
        )
        return list(table_fields_set & set(fields))

class Model(CRUDMixin, db.Model):
    """Base model class that includes CRUD convenience methods."""

    __abstract__ = True

    status_remove = -1
    status_default = 0
    status_available = 1

# From Mike Bayer's "Building the app" talk
# https://speakerdeck.com/zzzeek/building-the-app
class SurrogatePK(object):
    """A mixin that adds a surrogate integer 'primary key' column named ``id`` to any declarative-mapped class."""

    __table_args__ = {'extend_existing': True}

    id = db.Column(db.Integer, primary_key=True)

    @classmethod
    def get_by_id(cls, record_id):
        """Get record by ID."""
        if any(
                (isinstance(record_id, basestring) and record_id.isdigit(),
                 isinstance(record_id, (int))),
        ):
            return cls.query.get(int(record_id))
        return None

def reference_col(tablename, nullable=False, pk_name='id', **kwargs):
    """Column that adds primary key foreign key reference.

    Usage: ::

        category_id = reference_col('category')
        category = relationship('Category', backref='categories')
    """
    return db.Column(
            db.ForeignKey('{0}.{1}'.format(tablename, pk_name)),
            nullable=nullable, **kwargs)
<commit_msg>Add flake8: noqa linter directive to database.py<commit_after># -*- coding: utf-8 -*-
"""Database module, including the SQLAlchemy database object and DB-related utilities."""

# flake8: noqa  # flake8 has real problems linting this file on Python 2

from pprint import pformat

from sqlalchemy import desc, or_
from sqlalchemy.sql.sqltypes import Date, DateTime
from werkzeug import cached_property
from flask import current_app
from walle.service.extensions import db
from walle.service.utils import basestring
from walle.service.utils import datetime_str_to_obj, date_str_to_obj

# Alias common SQLAlchemy names
Column = db.Column
relationship = db.relationship

OPERATOR_FUNC_DICT = {
    '=': (lambda cls, k, v: getattr(cls, k) == v),
    '==': (lambda cls, k, v: getattr(cls, k) == v),
    'eq': (lambda cls, k, v: getattr(cls, k) == v),
    '!=': (lambda cls, k, v: getattr(cls, k) != v),
    'ne': (lambda cls, k, v: getattr(cls, k) != v),
    'neq': (lambda cls, k, v: getattr(cls, k) != v),
    '>': (lambda cls, k, v: getattr(cls, k) > v),
    'gt': (lambda cls, k, v: getattr(cls, k) > v),
    '>=': (lambda cls, k, v: getattr(cls, k) >= v),
    'gte': (lambda cls, k, v: getattr(cls, k) >= v),
    '<': (lambda cls, k, v: getattr(cls, k) < v),
    'lt': (lambda cls, k, v: getattr(cls, k) < v),
    '<=': (lambda cls, k, v: getattr(cls, k) <= v),
    'lte': (lambda cls, k, v: getattr(cls, k) <= v),
    'or': (lambda cls, k, v: or_(getattr(cls, k) == value for value in v)),
    'in': (lambda cls, k, v: getattr(cls, k).in_(v)),
    'nin': (lambda cls, k, v: ~getattr(cls, k).in_(v)),
    'like': (lambda cls, k, v: getattr(cls, k).like('%%%s%%' % (v))),
    'nlike': (lambda cls, k, v: ~getattr(cls, k).like(v)),
    '+': (lambda cls, k, v: getattr(cls, k) + v),
    'incr': (lambda cls, k, v: getattr(cls, k) + v),
    '-': (lambda cls, k, v: getattr(cls, k) - v),
    'decr': (lambda cls, k, v: getattr(cls, k) - v),
}

def parse_operator(cls, filter_name_dict):
    """ 用来返回sqlalchemy query对象filter使用的表达式
    Args:
        filter_name_dict (dict): 过滤条件dict
        {
            'last_name': {'eq': 'wang'},    # 如果是dic使用key作为操作符
            'age': {'>': 12}
        }
    Returns:
        binary_expression_list (lambda list)
    """

    def _change_type(cls, field, value):
        """ 有些表字段比如DateTime类型比较的时候需要转换类型，
        前端传过来的都是字符串，Date等类型没法直接相比较，需要转成Date类型
        Args:
            cls (class): Model class
            field (str): Model class field
            value (str): value need to compare
        """
        field_type = getattr(cls, field).type
        if isinstance(field_type, Date):
            return date_str_to_obj(value)
        elif isinstance(field_type, DateTime):
            return datetime_str_to_obj(value)
        else:
            return value

    binary_expression_list = []
    for field, op_dict in list(filter_name_dict.items()):
        for op, op_val in list(op_dict.items()):
            op_val = _change_type(cls, field, op_val)
            if op in OPERATOR_FUNC_DICT:
                binary_expression_list.append(
                        OPERATOR_FUNC_DICT[op](cls, field, op_val)
                )
    return binary_expression_list

class CRUDMixin(object):
    """Mixin that adds convenience methods for
    CRUD (create, read, update, delete) operations."""

    @classmethod
    def create(cls, **kwargs):
        """Create a new record and save it the database."""
        instance = cls(**kwargs)
        return instance.save()

    @classmethod
    def create_from_dict(cls, d):
        """Create a new record and save it the database."""
        assert isinstance(d, dict)
        instance = cls(**d)
        return instance.save()

    def update(self, commit=True, **kwargs):
        """Update specific fields of a record."""
        for attr, value in list(kwargs.items()):
            setattr(self, attr, value)
        return commit and self.save() or self

    def save(self, commit=True):
        """Save the record."""
        db.session.add(self)
        if commit:
            try:
                db.session.commit()
            except Exception as e:
                current_app.logger.info(e)
                db.session.rollback()
        return self

    def delete(self, commit=True):
        """Remove the record from the database."""
        db.session.delete(self)
        if commit:
            try:
                db.session.commit()
            except:
                db.session.rollback()
        return self

    def to_dict(self, fields_list=None):
        """
        Args:
            fields (str list): 指定返回的字段
        """
        column_list = fields_list or [
            column.name for column in self.__table__.columns
            ]
        return {
            column_name: getattr(self, column_name)
            for column_name in column_list
            }

    @classmethod
    def create_or_update(cls, query_dict, update_dict=None):
        instance = db.session.query(cls).filter_by(**query_dict).first()
        if instance:  # update
            if update_dict is not None:
                return instance.update(**update_dict)
            else:
                return instance
        else:  # create new instance
            query_dict.update(update_dict or {})
            return cls.create(**query_dict)

    @classmethod
    def query_paginate(cls, page=1, limit=20, fields=None, order_by_list=[('id', 'desc')],
                       filter_name_dict=None):
        """ 通用的分页查询函数
        Args:
            page (int): 页数
            limit (int): 每页个数
            order_by_list (tuple list): 用来指定排序的字段，可以传多个
                [ ('id', 1), ('name', -1) ]，1表示正序，-1表示逆序
                or
                [ ('id', 'asc'), ('name', 'desc') ]，1表示正序，-1表示逆序

            filter_name_dict (dict): 过滤条件，使用字典表示，使用字段名作为key，value
                是{'operator': to_compare_value}, e.g.:
                {
                    'last_name': {'eq': 'wang'},  # 如果是dic使用key作为操作符
                    'age': {'>': 12}
                }

        Returns:
            if fields is not None:
                (keytuple_list, total_cnt) (tuple)
            else:
                (instance_list, total_cnt) (tuple)

        前段查询参数规范：
        request.args 示例：
        ImmutableMultiDict([('limit', '10'), ('page', '1'), ('filter', '[{"field":"name","op":"eq","q":"g"},{
        "field":"id","op":"gt","q":"5"
        }]')])

        page: 页码
        limit: 每页限制
        order: 顺序，取值"asc" "desc". """'name', 'asc', 'model', 'desc'"""
        fields: 需要返回的字段
        filter: 过滤条件：
        {
            field: 需要过滤的字段
            op: 过滤操作符,支持"eq","neq","gt","gte","lt","lte","in","nin","like"
            q: 过滤值
        }
        """
        fields = (
            [getattr(cls, column) for column in fields] if fields is not None
            else None
        )
        if fields:
            query = db.session.query(*fields)
        else:
            query = db.session.query(cls)
        if order_by_list:
            for (field, order) in order_by_list:
                query = (
                    query.order_by(getattr(cls, field)) if order == 1 else
                    query.order_by(desc(getattr(cls, field)))
                )

        if filter_name_dict:
            p = parse_operator(cls, filter_name_dict)
            query = query.filter(*p)

        total_cnt = query.count()
        start = (page - 1) * limit
        query = query.offset(start).limit(limit)
        instance_or_keytuple_list = query.all()
        return instance_or_keytuple_list, total_cnt

    @classmethod
    def dump_schema(cls, items, fields, schema_class):
        """ 用来序列化从数据库查询出来的对象
        Args:
            items (instance list): obj list query from mysql
            fields (str list): fields need to dump
            schema_class (marshmallow.Schema): marshmallow.Schema class
        Returns:
            items, err
        """
        fields = (
            fields if fields else list(schema_class._declared_fields.keys())
        )
        schema = schema_class(many=True, only=fields)
        items, err = schema.dump(items)
        return items, err

    @classmethod
    def query_paginate_and_dump_schema(cls, page=1, limit=20, fields=None,
                                       order_by_list=None,
                                       filter_name_dict=None,
                                       schema_class=None):
        """ 分页查询并且返回dump后的对象，可以解决大部分查询问题 """
        assert schema_class
        items, total_cnt = cls.query_paginate(
                page, limit, fields, order_by_list, filter_name_dict
        )
        items, err = cls.dump_schema(items, fields, schema_class)
        return items, total_cnt

    def __repr__(self):
        return pformat(self.to_dict())

    @cached_property
    def column_name_set(self):
        return set([column.name for column in self.__table__.columns])

    @classmethod
    def get_common_fields(cls, fields=None):
        """ 防止传过来的fields含有该Model没有的字段 """
        if not fields:
            return []
        table_fields_set = set(
                [column.name for column in cls.__table__.columns]
        )
        return list(table_fields_set & set(fields))

class Model(CRUDMixin, db.Model):
    """Base model class that includes CRUD convenience methods."""

    __abstract__ = True

    status_remove = -1
    status_default = 0
    status_available = 1

# From Mike Bayer's "Building the app" talk
# https://speakerdeck.com/zzzeek/building-the-app
class SurrogatePK(object):
    """A mixin that adds a surrogate integer 'primary key' column named ``id`` to any declarative-mapped class."""

    __table_args__ = {'extend_existing': True}

    id = db.Column(db.Integer, primary_key=True)

    @classmethod
    def get_by_id(cls, record_id):
        """Get record by ID."""
        if any(
                (isinstance(record_id, basestring) and record_id.isdigit(),
                 isinstance(record_id, (int))),
        ):
            return cls.query.get(int(record_id))
        return None

def reference_col(tablename, nullable=False, pk_name='id', **kwargs):
    """Column that adds primary key foreign key reference.

    Usage: ::

        category_id = reference_col('category')
        category = relationship('Category', backref='categories')
    """
    return db.Column(
            db.ForeignKey('{0}.{1}'.format(tablename, pk_name)),
            nullable=nullable, **kwargs)
<|endoftext|>


<commit_before># -*- coding: utf-8 -*-
__version__ = '1.0.0'
<commit_msg>Bumped version to 1.0.1<commit_after># -*- coding: utf-8 -*-
__version__ = '1.0.1'
<|endoftext|>


<commit_before>from genotype import Genotype
from phenotype import Phenotype
from ped_writer import PedWriter
import os
import subprocess

class EpactsModel(object):
    __models = []

    def __init__(self, config, cmd=""):
        self.cmd = cmd
        self.config = config
        self.cores_per_job = 56

    def get_opts(self, model, geno):
        opts = []
        if model.get("response_invnorm", False):
            opts.append("--inv-norm")
        return opts

    def get_ped_writer(self, model, pheno):
        ped_writer = PedWriter(pheno.get_pheno_reader(), \
            model["response"], model.get("covariates",[]))
        return ped_writer

    def get_analysis_command(self, model, geno, ped_path, ped_writer):
        cmd = "{} {}".format(self.config.get("ANALYSIS_BINARY", "epacts"), self.cmd) + \
            " --vcf {}".format(geno.get_vcf_path(1)) + \
            " --ped {}".format(ped_path) +  \
            " --field GT" + \
            " --sepchr" + \
            " --out ./output --run {}".format(self.cores_per_job)
        for resp in ped_writer.get_response_headers():
            cmd += " --pheno {}".format(resp)
        for covar in ped_writer.get_covar_headers():
            cmd += " --cov {}".format(covar)
        cmd += " " + " ".join(self.get_opts(model, geno))
        return cmd

    def get_postprocessing_command(self):
        cmds = []
        if self.cmd == "group":
            cmds.append("if [ -e output.epacts -a ! -e output.epacts.gz ]; then\n" + \
                "  awk 'NR<2{print;next}{print| \"sort -g -k1,1 -k2g,3\"}' output.epacts | " + \
                "{} -c > output.epacts.gz\n".format(self.config.get("BGZIP_BINARY", "bgzip")) + \
                " {} -p bed output.epacts.gz\n".format(self.config.get("TABIX_BINARY", "tabix")) + \
                "fi")
        if self.config.get("MANHATTAN_BINARY"):
            cmd  = "{} ./output.epacts.gz ./manhattan.json".format(self.config.get("MANHATTAN_BINARY", ""))
            cmds.append(cmd)
        if self.config.get("TOPHITS_BINARY"):
            cmd = "{} ./output.epacts.gz ./qq.json".format(self.config.get("QQPLOT_BINARY", ""))
            cmds.append(cmd)
        if self.config.get("TOPHITS_BINARY"):
            cmd =  "{} ./output.epacts.top5000 ./tophits.json".format(self.config.get("TOPHITS_BINARY"))
            if self.cmd == "group":
                cmd += " --window 0"
            elif self.config.get("NEAREST_GENE_BED"):
                cmd += " --gene " + self.config.get("NEAREST_GENE_BED")
            cmds.append(cmd)
        return "\n".join(cmds)

    @staticmethod
    def get(model, config):
        for m in EpactsModel.__models:
            if m.model_code == model:
                return m(config)
        raise ValueError('Unrecognized model type: %s' % (model_type,))

    @staticmethod
    def list():
        def desc(x):
            return {"code": x.model_code, "name": x.model_name, "description": x.model_desc}
        return [ desc(m) for m in EpactsModel.__models]

    @staticmethod
    def register(m):
        EpactsModel.__models.append(m)

class LMEpactsModel(EpactsModel):
    model_code = "lm"
    model_name = "Linear Wald Test"
    model_desc = "A simple linear model"

    def __init__(self, config):
        EpactsModel.__init__(self, config, "single")

    def get_opts(self, model, geno):
        opts = super(self.__class__, self).get_opts(model, geno)
        opts += ["--test q.linear",
            "--unit 500000",
            "--min-maf 0.001" ]
        return opts
EpactsModel.register(LMEpactsModel)

class LMMEpactsModel(EpactsModel):
    model_code = "lmm"
    model_name = "Linear Mixed Model"
    model_desc = "Adjust for potential relatedness using kinship matrix"

    def __init__(self, config):
        EpactsModel.__init__(self, config, "single")

    def get_opts(self, model, geno):
        opts = super(self.__class__, self).get_opts(model, geno)
        opts += ["--test q.emmax",
            "--kin {}".format(geno.get_kinship_path()),
            "--unit 500000",
            "--min-maf 0.001"]
        return opts
EpactsModel.register(LMMEpactsModel)

class SkatOEpactsModel(EpactsModel):
    model_code = "skato"
    model_name = "SKAT-O Test"
    model_desc = "Adaptive burden test"

    def __init__(self, config):
        EpactsModel.__init__(self, config, "group")

    def get_opts(self, model, geno):
        opts = super(self.__class__, self).get_opts(model, geno)
        group = model.get("group", "nonsyn")
        opts += ["--test skat",
            "--skat-o",
            "--groupf {}".format(geno.get_groups_path(group)),
            "--unit 500",
            "--max-maf 0.05"]
        return opts
EpactsModel.register(SkatOEpactsModel)

class MMSkatOEpactsModel(EpactsModel):
    model_code = "mmskato"
    model_name = "Mixed Model SKAT-O Test"
    model_desc = "Adaptive burden test that adjusts for potential relatedness using kinship matrix"

    def __init__(self, config):
        EpactsModel.__init__(self, config, "group")

    def get_opts(self, model, geno):
        opts = super(self.__class__, self).get_opts(model, geno)
        group = model.get("group", "nonsyn")
        opts += ["--test mmskat",
            "--skat-o",
            "--groupf {}".format(geno.get_groups_path(group)),
            "--kin {}".format(geno.get_kinship_path()),
            "--unit 500",
            "--max-maf 0.05"]
        return opts
EpactsModel.register(MMSkatOEpactsModel)

class SlurmEpactsJob:

    def __init__(self, job_id, job_directory, config = None):
        self.job_id = job_id
        self.job_directory = job_directory
        if config:
            self.config = config
        else:
            self.config = dict()

    def create_sbatch_header(self, epm):
       return "#!/bin/bash\n" + \
           "#SBATCH --partition=encore\n" + \
           "#SBATCH --job-name=gasp_{}\n".format(self.job_id)  + \
           "#SBATCH --cpus-per-task={}\n".format(epm.cores_per_job)  + \
           "#SBATCH --workdir={}\n".format(self.job_directory) + \
           "#SBATCH --mem-per-cpu=4000\n" + \
           "#SBATCH --time=14-0\n" + \
           "#SBATCH --nodes=1\n"

    def create_analysis_command(self, epm, job_desc, geno=None, pheno=None):
        if geno is None:
            if "genotype" in job_desc:
                geno = Genotype.get(job_desc["genotype"], self.config)
            else:
                raise Exception("No genotype information in job")

        if pheno is None:
            if "phenotype" in job_desc:
                pheno = Phenotype.get(job_desc["phenotype"], self.config)
            else:
                raise Exception("No phenotype information in job")

        ped_writer = epm.get_ped_writer(job_desc, pheno)
        ped_path = self.relative_path("pheno.ped")
        with open(ped_path,"w") as pedfile:
            ped_writer.write_to_file(pedfile)

        return epm.get_analysis_command(job_desc, geno, ped_path, ped_writer)

    def create_postprocessing_command(self, epm):
        cmd =  "if [ $EXIT_STATUS == 0 ]; then\n"
        cmd += epm.get_postprocessing_command()
        cmd += "\nfi\n"
        return cmd

    def create_launch_script(self, job_desc):
        model_type = job_desc.get("type", None)
        epm = EpactsModel.get(model_type, self.config)

        cmd = self.create_sbatch_header(epm) + "\n"
        cmd += "{} 2> ./err.log 1> ./out.log\n".format(self.create_analysis_command(epm, job_desc))
        cmd += "EXIT_STATUS=$?\n"
        cmd += self.create_postprocessing_command(epm) + "\n"
        cmd += "echo $EXIT_STATUS > ./exit_status.txt\n"
        cmd += "exit $EXIT_STATUS\n"
        return cmd

    def submit_job(self, job_desc):
        sbatch = self.config.get("QUEUE_JOB_BINARY", "sbatch")
        batch_script_path = self.relative_path("batch_script.sh")
        batch_output_path = self.relative_path("batch_script_output.txt")
        with open(batch_script_path, "w") as f:
            f.write(self.create_launch_script(job_desc))
        with open(batch_output_path, "w") as f:
            try:
                subprocess.call([sbatch, batch_script_path], stdout=f)
            except subprocess.CalledProcessError:
                raise Exception("Could not queue job")
            except OSError:
                raise Exception("Could not find sbatch")
        return True

    def relative_path(self, *args):
        return os.path.expanduser(os.path.join(self.job_directory, *args))

    @staticmethod
    def available_models():
        return EpactsModel.list()
<commit_msg>Increase memory per job<commit_after>from genotype import Genotype
from phenotype import Phenotype
from ped_writer import PedWriter
import os
import subprocess

class EpactsModel(object):
    __models = []

    def __init__(self, config, cmd=""):
        self.cmd = cmd
        self.config = config
        self.cores_per_job = 56

    def get_opts(self, model, geno):
        opts = []
        if model.get("response_invnorm", False):
            opts.append("--inv-norm")
        return opts

    def get_ped_writer(self, model, pheno):
        ped_writer = PedWriter(pheno.get_pheno_reader(), \
            model["response"], model.get("covariates",[]))
        return ped_writer

    def get_analysis_command(self, model, geno, ped_path, ped_writer):
        cmd = "{} {}".format(self.config.get("ANALYSIS_BINARY", "epacts"), self.cmd) + \
            " --vcf {}".format(geno.get_vcf_path(1)) + \
            " --ped {}".format(ped_path) +  \
            " --field GT" + \
            " --sepchr" + \
            " --out ./output --run {}".format(self.cores_per_job)
        for resp in ped_writer.get_response_headers():
            cmd += " --pheno {}".format(resp)
        for covar in ped_writer.get_covar_headers():
            cmd += " --cov {}".format(covar)
        cmd += " " + " ".join(self.get_opts(model, geno))
        return cmd

    def get_postprocessing_command(self):
        cmds = []
        if self.cmd == "group":
            cmds.append("if [ -e output.epacts -a ! -e output.epacts.gz ]; then\n" + \
                "  awk 'NR<2{print;next}{print| \"sort -g -k1,1 -k2g,3\"}' output.epacts | " + \
                "{} -c > output.epacts.gz\n".format(self.config.get("BGZIP_BINARY", "bgzip")) + \
                " {} -p bed output.epacts.gz\n".format(self.config.get("TABIX_BINARY", "tabix")) + \
                "fi")
        if self.config.get("MANHATTAN_BINARY"):
            cmd  = "{} ./output.epacts.gz ./manhattan.json".format(self.config.get("MANHATTAN_BINARY", ""))
            cmds.append(cmd)
        if self.config.get("TOPHITS_BINARY"):
            cmd = "{} ./output.epacts.gz ./qq.json".format(self.config.get("QQPLOT_BINARY", ""))
            cmds.append(cmd)
        if self.config.get("TOPHITS_BINARY"):
            cmd =  "{} ./output.epacts.top5000 ./tophits.json".format(self.config.get("TOPHITS_BINARY"))
            if self.cmd == "group":
                cmd += " --window 0"
            elif self.config.get("NEAREST_GENE_BED"):
                cmd += " --gene " + self.config.get("NEAREST_GENE_BED")
            cmds.append(cmd)
        return "\n".join(cmds)

    @staticmethod
    def get(model, config):
        for m in EpactsModel.__models:
            if m.model_code == model:
                return m(config)
        raise ValueError('Unrecognized model type: %s' % (model_type,))

    @staticmethod
    def list():
        def desc(x):
            return {"code": x.model_code, "name": x.model_name, "description": x.model_desc}
        return [ desc(m) for m in EpactsModel.__models]

    @staticmethod
    def register(m):
        EpactsModel.__models.append(m)

class LMEpactsModel(EpactsModel):
    model_code = "lm"
    model_name = "Linear Wald Test"
    model_desc = "A simple linear model"

    def __init__(self, config):
        EpactsModel.__init__(self, config, "single")

    def get_opts(self, model, geno):
        opts = super(self.__class__, self).get_opts(model, geno)
        opts += ["--test q.linear",
            "--unit 500000",
            "--min-maf 0.001" ]
        return opts
EpactsModel.register(LMEpactsModel)

class LMMEpactsModel(EpactsModel):
    model_code = "lmm"
    model_name = "Linear Mixed Model"
    model_desc = "Adjust for potential relatedness using kinship matrix"

    def __init__(self, config):
        EpactsModel.__init__(self, config, "single")

    def get_opts(self, model, geno):
        opts = super(self.__class__, self).get_opts(model, geno)
        opts += ["--test q.emmax",
            "--kin {}".format(geno.get_kinship_path()),
            "--unit 500000",
            "--min-maf 0.001"]
        return opts
EpactsModel.register(LMMEpactsModel)

class SkatOEpactsModel(EpactsModel):
    model_code = "skato"
    model_name = "SKAT-O Test"
    model_desc = "Adaptive burden test"

    def __init__(self, config):
        EpactsModel.__init__(self, config, "group")

    def get_opts(self, model, geno):
        opts = super(self.__class__, self).get_opts(model, geno)
        group = model.get("group", "nonsyn")
        opts += ["--test skat",
            "--skat-o",
            "--groupf {}".format(geno.get_groups_path(group)),
            "--unit 500",
            "--max-maf 0.05"]
        return opts
EpactsModel.register(SkatOEpactsModel)

class MMSkatOEpactsModel(EpactsModel):
    model_code = "mmskato"
    model_name = "Mixed Model SKAT-O Test"
    model_desc = "Adaptive burden test that adjusts for potential relatedness using kinship matrix"

    def __init__(self, config):
        EpactsModel.__init__(self, config, "group")

    def get_opts(self, model, geno):
        opts = super(self.__class__, self).get_opts(model, geno)
        group = model.get("group", "nonsyn")
        opts += ["--test mmskat",
            "--skat-o",
            "--groupf {}".format(geno.get_groups_path(group)),
            "--kin {}".format(geno.get_kinship_path()),
            "--unit 500",
            "--max-maf 0.05"]
        return opts
EpactsModel.register(MMSkatOEpactsModel)

class SlurmEpactsJob:

    def __init__(self, job_id, job_directory, config = None):
        self.job_id = job_id
        self.job_directory = job_directory
        if config:
            self.config = config
        else:
            self.config = dict()

    def create_sbatch_header(self, epm):
       return "#!/bin/bash\n" + \
           "#SBATCH --partition=encore\n" + \
           "#SBATCH --job-name=gasp_{}\n".format(self.job_id)  + \
           "#SBATCH --cpus-per-task={}\n".format(epm.cores_per_job)  + \
           "#SBATCH --workdir={}\n".format(self.job_directory) + \
           "#SBATCH --mem-per-cpu=7000\n" + \
           "#SBATCH --time=14-0\n" + \
           "#SBATCH --nodes=1\n"

    def create_analysis_command(self, epm, job_desc, geno=None, pheno=None):
        if geno is None:
            if "genotype" in job_desc:
                geno = Genotype.get(job_desc["genotype"], self.config)
            else:
                raise Exception("No genotype information in job")

        if pheno is None:
            if "phenotype" in job_desc:
                pheno = Phenotype.get(job_desc["phenotype"], self.config)
            else:
                raise Exception("No phenotype information in job")

        ped_writer = epm.get_ped_writer(job_desc, pheno)
        ped_path = self.relative_path("pheno.ped")
        with open(ped_path,"w") as pedfile:
            ped_writer.write_to_file(pedfile)

        return epm.get_analysis_command(job_desc, geno, ped_path, ped_writer)

    def create_postprocessing_command(self, epm):
        cmd =  "if [ $EXIT_STATUS == 0 ]; then\n"
        cmd += epm.get_postprocessing_command()
        cmd += "\nfi\n"
        return cmd

    def create_launch_script(self, job_desc):
        model_type = job_desc.get("type", None)
        epm = EpactsModel.get(model_type, self.config)

        cmd = self.create_sbatch_header(epm) + "\n"
        cmd += "{} 2> ./err.log 1> ./out.log\n".format(self.create_analysis_command(epm, job_desc))
        cmd += "EXIT_STATUS=$?\n"
        cmd += self.create_postprocessing_command(epm) + "\n"
        cmd += "echo $EXIT_STATUS > ./exit_status.txt\n"
        cmd += "exit $EXIT_STATUS\n"
        return cmd

    def submit_job(self, job_desc):
        sbatch = self.config.get("QUEUE_JOB_BINARY", "sbatch")
        batch_script_path = self.relative_path("batch_script.sh")
        batch_output_path = self.relative_path("batch_script_output.txt")
        with open(batch_script_path, "w") as f:
            f.write(self.create_launch_script(job_desc))
        with open(batch_output_path, "w") as f:
            try:
                subprocess.call([sbatch, batch_script_path], stdout=f)
            except subprocess.CalledProcessError:
                raise Exception("Could not queue job")
            except OSError:
                raise Exception("Could not find sbatch")
        return True

    def relative_path(self, *args):
        return os.path.expanduser(os.path.join(self.job_directory, *args))

    @staticmethod
    def available_models():
        return EpactsModel.list()
<|endoftext|>


<commit_before>from django.test import TestCase, Client
from django.urls import reverse
from web.views import index, about, compare, reference

class TestViews(TestCase):

	def test_index_view_GET(self):
		url = reverse('index')
		response = self.client.get(url)

		self.assertEqual(response.status_code, 200)
		self.assertTemplateUsed(response, 'index.html')
		self.assertTemplateUsed(response, 'base.html')

	def test_about_view_GET(self):
		url = reverse('about')
		response = self.client.get(url)

		self.assertEqual(response.status_code, 200)
		self.assertTemplateUsed(response, 'about.html')
		self.assertTemplateUsed(response, 'base.html')

	def test_compare_view_both_valid_languages(self):
		url = reverse('compare') + '?concept=data_types&lang1=python&lang2=java'
		response = self.client.get(url)

		self.assertEquals(response.status_code, 200)
		self.assertTemplateUsed(response, 'compare.html')
		self.assertTemplateUsed(response, 'base.html')

	def test_compare_view_invalid_languages(self):
		url = reverse('compare') + '?concept=data_types&lang1=cupcake&lang2=donut'
		response = self.client.get(url)

		self.assertEquals(response.status_code, 404)
		self.assertTemplateNotUsed(response, 'compare.html')
		self.assertTemplateNotUsed(response, 'base.html')

	def test_compare_view_one_valid_one_invalid_language(self):
		url = reverse('compare') + '?concept=data_types&lang1=python&lang2=donut'
		response = self.client.get(url)

		self.assertEquals(response.status_code, 404)
		self.assertTemplateNotUsed(response, 'compare.html')
		self.assertTemplateNotUsed(response, 'base.html')

	def test_reference_view_valid_language(self):
		url = reverse('reference') + '?concept=data_types&lang=python'
		response = self.client.get(url)

		self.assertEquals(response.status_code, 200)
		self.assertTemplateUsed(response, 'reference.html')
		self.assertTemplateUsed(response, 'base.html')

	def test_reference_view_invalid_languages(self):
		url = reverse('reference') + '?concept=data_types&lang=cupcake'
		response = self.client.get(url)

		self.assertEquals(response.status_code, 404)
		self.assertTemplateNotUsed(response, 'reference.html')
		self.assertTemplateNotUsed(response, 'base.html')

<commit_msg>Add tests for compare/reference pages for testing empty query strings<commit_after>from django.test import TestCase, Client
from django.urls import reverse
from web.views import index, about, compare, reference

class TestViews(TestCase):

	def test_index_view_GET(self):
		url = reverse('index')
		response = self.client.get(url)

		self.assertEqual(response.status_code, 200)
		self.assertTemplateUsed(response, 'index.html')
		self.assertTemplateUsed(response, 'base.html')

	def test_about_view_GET(self):
		url = reverse('about')
		response = self.client.get(url)

		self.assertEqual(response.status_code, 200)
		self.assertTemplateUsed(response, 'about.html')
		self.assertTemplateUsed(response, 'base.html')

	def test_compare_view_both_valid_languages(self):
		url = reverse('compare') + '?concept=data_types&lang1=python&lang2=java'
		response = self.client.get(url)

		self.assertEquals(response.status_code, 200)
		self.assertTemplateUsed(response, 'compare.html')
		self.assertTemplateUsed(response, 'base.html')

	def test_compare_view_invalid_languages(self):
		url = reverse('compare') + '?concept=data_types&lang1=cupcake&lang2=donut'
		response = self.client.get(url)

		self.assertEquals(response.status_code, 404)
		self.assertTemplateNotUsed(response, 'compare.html')
		self.assertTemplateNotUsed(response, 'base.html')

	def test_compare_view_one_valid_one_invalid_language(self):
		url = reverse('compare') + '?concept=data_types&lang1=python&lang2=donut'
		response = self.client.get(url)

		self.assertEquals(response.status_code, 404)
		self.assertTemplateNotUsed(response, 'compare.html')
		self.assertTemplateNotUsed(response, 'base.html')

	def test_compare_view_empty_query_string(self):
		url = reverse('compare')

		# This SHOULD be a pretty 404 error, but presently is a Python exception
		with self.assertRaises(ValueError):
			response = self.client.get(url)

		# After implemented, this SHOULD eventually be:
		# response = self.client.get(url)
		# self.assertEquals(response.status_code, 404)
		# self.assertTemplateNotUsed(response, 'some_new_error_template.html')

	def test_compare_view_invalid_concept(self):
		url = reverse('compare') + '?concept=boop'

		# This SHOULD be a pretty 404 error, but presently is a Python exception
		with self.assertRaises(ValueError):
			response = self.client.get(url)

		# After implemented, this SHOULD eventually be:
		# response = self.client.get(url)
		# self.assertEquals(response.status_code, 404)
		# self.assertTemplateNotUsed(response, 'some_new_error_template.html')

	def test_reference_view_valid_language(self):
		url = reverse('reference') + '?concept=data_types&lang=python'
		response = self.client.get(url)

		self.assertEquals(response.status_code, 200)
		self.assertTemplateUsed(response, 'reference.html')
		self.assertTemplateUsed(response, 'base.html')

	def test_reference_view_invalid_languages(self):
		url = reverse('reference') + '?concept=data_types&lang=cupcake'
		response = self.client.get(url)

		self.assertEquals(response.status_code, 404)
		self.assertTemplateNotUsed(response, 'reference.html')
		self.assertTemplateNotUsed(response, 'base.html')

	def test_reference_view_empty_query_string(self):
		url = reverse('reference')

		# This SHOULD be a pretty 404 error, but presently is a Python exception
		with self.assertRaises(ValueError):
			response = self.client.get(url)

		# After implemented, this SHOULD eventually be:
		# response = self.client.get(url)
		# self.assertEquals(response.status_code, 404)
		# self.assertTemplateNotUsed(response, 'some_new_error_template.html')

	def test_reference_view_invalid_concept(self):
		url = reverse('reference') + '?concept=boop'

		# This SHOULD be a pretty 404 error, but presently is a Python exception
		with self.assertRaises(ValueError):
			response = self.client.get(url)

		# After implemented, this SHOULD eventually be:
		# response = self.client.get(url)
		# self.assertEquals(response.status_code, 404)
		# self.assertTemplateNotUsed(response, 'some_new_error_template.html')
<|endoftext|>


<commit_before># -*- coding: utf-8 -*-
import time
import logging
import functools
import httplib as http

import bleach

from flask import request

from modularodm import Q
from framework.auth.decorators import collect_auth
from framework.auth.decorators import must_be_logged_in

from website.models import Node
from website.models import User
from website.search import util
from website.util import api_url_for
from website.search import exceptions
from website.search import share_search
import website.search.search as search
from framework.exceptions import HTTPError
from website.search.exceptions import IndexNotFoundError
from website.search.exceptions import MalformedQueryError
from website.search.util import build_query
from website.project.views.contributor import get_node_contributors_abbrev

logger = logging.getLogger(__name__)

RESULTS_PER_PAGE = 250

def handle_search_errors(func):
    @functools.wraps(func)
    def wrapped(*args, **kwargs):
        try:
            return func(*args, **kwargs)
        except exceptions.MalformedQueryError:
            raise HTTPError(http.BAD_REQUEST, data={
                'message_short': 'Bad search query',
                'message_long': ('Please check our help (the question mark beside the search box) for more information '
                                 'on advanced search queries.'),
            })
        except exceptions.SearchUnavailableError:
            raise HTTPError(http.SERVICE_UNAVAILABLE, data={
                'message_short': 'Search unavailable',
                'message_long': ('Our search service is currently unavailable, if the issue persists, '
                'please report it to support@osf.io.'),
            })
    return wrapped

@handle_search_errors
def search_search(**kwargs):
    _type = kwargs.get('type', None)

    tick = time.time()
    results = {}

    if request.method == 'POST':
        results = search.search(request.get_json(), doc_type=_type)
    elif request.method == 'GET':
        q = request.args.get('q', '*')
        # TODO Match javascript params?
        start = request.args.get('from', '0')
        size = request.args.get('size', '10')
        results = search.search(build_query(q, start, size), doc_type=_type)

    results['time'] = round(time.time() - tick, 2)
    return results

def conditionally_add_query_item(query, item, condition):
    """ Helper for the search_projects_by_title function which will add a condition to a query
    It will give an error if the proper search term is not used.
    :param query: The modular ODM query that you want to modify
    :param item:  the field to query on
    :param condition: yes, no, or either
    :return: the modified query
    """

    condition = condition.lower()

    if condition == "yes":
        return query & Q(item, 'eq', True)
    elif condition == "no":
        return query & Q(item, 'eq', False)
    elif condition == "either":
        return query

    raise HTTPError(http.BAD_REQUEST)

@must_be_logged_in
def search_projects_by_title(**kwargs):
    """ Search for nodes by title. Can pass in arguments from the URL to modify the search
    :arg term: The substring of the title.
    :arg category: Category of the node.
    :arg isDeleted: yes, no, or either. Either will not add a qualifier for that argument in the search.
    :arg isFolder: yes, no, or either. Either will not add a qualifier for that argument in the search.
    :arg isRegistration: yes, no, or either. Either will not add a qualifier for that argument in the search.
    :arg includePublic: yes or no. Whether the projects listed should include public projects.
    :arg includeContributed: yes or no. Whether the search should include projects the current user has
        contributed to.
    :arg ignoreNode: a list of nodes that should not be included in the search.
    :return: a list of dictionaries of projects

    """
    # TODO(fabianvf): At some point, it would be nice to do this with elastic search
    user = kwargs['auth'].user

    term = request.args.get('term', '')
    max_results = int(request.args.get('maxResults', '10'))
    category = request.args.get('category', 'project').lower()
    is_deleted = request.args.get('isDeleted', 'no').lower()
    is_folder = request.args.get('isFolder', 'no').lower()
    is_registration = request.args.get('isRegistration', 'no').lower()
    include_public = request.args.get('includePublic', 'yes').lower()
    include_contributed = request.args.get('includeContributed', 'yes').lower()
    ignore_nodes = request.args.getlist('ignoreNode', [])

    matching_title = (
        Q('title', 'icontains', term) &  # search term (case insensitive)
        Q('category', 'eq', category)  # is a project
    )

    matching_title = conditionally_add_query_item(matching_title, 'is_deleted', is_deleted)
    matching_title = conditionally_add_query_item(matching_title, 'is_folder', is_folder)
    matching_title = conditionally_add_query_item(matching_title, 'is_registration', is_registration)

    if len(ignore_nodes) > 0:
        for node_id in ignore_nodes:
            matching_title = matching_title & Q('_id', 'ne', node_id)

    my_projects = []
    my_project_count = 0
    public_projects = []

    if include_contributed == "yes":
        my_projects = Node.find(
            matching_title &
            Q('contributors', 'contains', user._id)  # user is a contributor
        ).limit(max_results)
        my_project_count = my_project_count

    if my_project_count < max_results and include_public == "yes":
        public_projects = Node.find(
            matching_title &
            Q('is_public', 'eq', True)  # is public
        ).limit(max_results - my_project_count)

    results = list(my_projects) + list(public_projects)
    ret = process_project_search_results(results, **kwargs)
    return ret

@must_be_logged_in
def process_project_search_results(results, **kwargs):
    """
    :param results: list of projects from the modular ODM search
    :return: we return the entire search result, which is a list of
    dictionaries. This includes the list of contributors.
    """
    user = kwargs['auth'].user

    ret = []

    for project in results:
        authors = get_node_contributors_abbrev(project=project, auth=kwargs['auth'])
        authors_html = ''
        for author in authors['contributors']:
            a = User.load(author['user_id'])
            authors_html += '%s' % (a.url, a.fullname)
            authors_html += author['separator'] + ' '
        authors_html += ' ' + authors['others_count']

        ret.append({
            'id': project._id,
            'label': project.title,
            'value': project.title,
            'category': 'My Projects' if user in project.contributors else 'Public Projects',
            'authors': authors_html,
        })

    return ret

@collect_auth
def search_contributor(auth):
    user = auth.user if auth else None
    nid = request.args.get('excludeNode')
    exclude = Node.load(nid).contributors if nid else []
    query = bleach.clean(request.args.get('query', ''), tags=[], strip=True)
    page = int(bleach.clean(request.args.get('page', '0'), tags=[], strip=True))
    size = int(bleach.clean(request.args.get('size', '5'), tags=[], strip=True))
    return search.search_contributor(query=query, page=page, size=size,
                                     exclude=exclude, current_user=user)

@handle_search_errors
def search_share():
    tick = time.time()
    results = {}

    count = request.args.get('count') is not None
    raw = request.args.get('raw') is not None
    version = request.args.get('v')
    index = 'share_v{}'.format(version) if version else 'share'

    if request.method == 'POST':
        query = request.get_json()
    elif request.method == 'GET':
        query = build_query(
            request.args.get('q', '*'),
            request.args.get('from'),
            request.args.get('size'),
            sort=request.args.get('sort')
        )

    if count:
        results = search.count_share(query, index=index)
    else:
        results = search.search_share(query, raw, index=index)

    results['time'] = round(time.time() - tick, 2)
    return results

@handle_search_errors
def search_share_stats():
    q = request.args.get('q')
    query = build_query(q, 0, 0) if q else {}

    return search.share_stats(query=query)

@handle_search_errors
def search_share_atom(**kwargs):
    q = request.args.get('q', '*')
    sort = request.args.get('sort', 'dateUpdated')

    # we want the results per page to be constant between pages
    # TODO -  move this functionality into build_query in util
    start = util.compute_start(request.args.get('page', 1), RESULTS_PER_PAGE)

    query = build_query(q, size=RESULTS_PER_PAGE, start=start, sort=sort)

    try:
        search_results = search.search_share(query)
    except MalformedQueryError:
        raise HTTPError(http.BAD_REQUEST)
    except IndexNotFoundError:
        search_results = {
            'count': 0,
            'results': []
        }

    atom_url = api_url_for('search_share_atom', _xml=True, _absolute=True)

    return util.create_atom_feed(
        name='SHARE',
        data=search_results['results'],
        query=q,
        size=RESULTS_PER_PAGE,
        start=start,
        url=atom_url,
        to_atom=share_search.to_atom
    )

def search_share_providers():
    return search.share_providers()
<commit_msg>Atom also only queries with old metadata schema<commit_after># -*- coding: utf-8 -*-
import time
import logging
import functools
import httplib as http

import bleach

from flask import request

from modularodm import Q
from framework.auth.decorators import collect_auth
from framework.auth.decorators import must_be_logged_in

from website.models import Node
from website.models import User
from website.search import util
from website.util import api_url_for
from website.search import exceptions
from website.search import share_search
import website.search.search as search
from framework.exceptions import HTTPError
from website.search.exceptions import IndexNotFoundError
from website.search.exceptions import MalformedQueryError
from website.search.util import build_query
from website.project.views.contributor import get_node_contributors_abbrev

logger = logging.getLogger(__name__)

RESULTS_PER_PAGE = 250

def handle_search_errors(func):
    @functools.wraps(func)
    def wrapped(*args, **kwargs):
        try:
            return func(*args, **kwargs)
        except exceptions.MalformedQueryError:
            raise HTTPError(http.BAD_REQUEST, data={
                'message_short': 'Bad search query',
                'message_long': ('Please check our help (the question mark beside the search box) for more information '
                                 'on advanced search queries.'),
            })
        except exceptions.SearchUnavailableError:
            raise HTTPError(http.SERVICE_UNAVAILABLE, data={
                'message_short': 'Search unavailable',
                'message_long': ('Our search service is currently unavailable, if the issue persists, '
                'please report it to support@osf.io.'),
            })
    return wrapped

@handle_search_errors
def search_search(**kwargs):
    _type = kwargs.get('type', None)

    tick = time.time()
    results = {}

    if request.method == 'POST':
        results = search.search(request.get_json(), doc_type=_type)
    elif request.method == 'GET':
        q = request.args.get('q', '*')
        # TODO Match javascript params?
        start = request.args.get('from', '0')
        size = request.args.get('size', '10')
        results = search.search(build_query(q, start, size), doc_type=_type)

    results['time'] = round(time.time() - tick, 2)
    return results

def conditionally_add_query_item(query, item, condition):
    """ Helper for the search_projects_by_title function which will add a condition to a query
    It will give an error if the proper search term is not used.
    :param query: The modular ODM query that you want to modify
    :param item:  the field to query on
    :param condition: yes, no, or either
    :return: the modified query
    """

    condition = condition.lower()

    if condition == "yes":
        return query & Q(item, 'eq', True)
    elif condition == "no":
        return query & Q(item, 'eq', False)
    elif condition == "either":
        return query

    raise HTTPError(http.BAD_REQUEST)

@must_be_logged_in
def search_projects_by_title(**kwargs):
    """ Search for nodes by title. Can pass in arguments from the URL to modify the search
    :arg term: The substring of the title.
    :arg category: Category of the node.
    :arg isDeleted: yes, no, or either. Either will not add a qualifier for that argument in the search.
    :arg isFolder: yes, no, or either. Either will not add a qualifier for that argument in the search.
    :arg isRegistration: yes, no, or either. Either will not add a qualifier for that argument in the search.
    :arg includePublic: yes or no. Whether the projects listed should include public projects.
    :arg includeContributed: yes or no. Whether the search should include projects the current user has
        contributed to.
    :arg ignoreNode: a list of nodes that should not be included in the search.
    :return: a list of dictionaries of projects

    """
    # TODO(fabianvf): At some point, it would be nice to do this with elastic search
    user = kwargs['auth'].user

    term = request.args.get('term', '')
    max_results = int(request.args.get('maxResults', '10'))
    category = request.args.get('category', 'project').lower()
    is_deleted = request.args.get('isDeleted', 'no').lower()
    is_folder = request.args.get('isFolder', 'no').lower()
    is_registration = request.args.get('isRegistration', 'no').lower()
    include_public = request.args.get('includePublic', 'yes').lower()
    include_contributed = request.args.get('includeContributed', 'yes').lower()
    ignore_nodes = request.args.getlist('ignoreNode', [])

    matching_title = (
        Q('title', 'icontains', term) &  # search term (case insensitive)
        Q('category', 'eq', category)  # is a project
    )

    matching_title = conditionally_add_query_item(matching_title, 'is_deleted', is_deleted)
    matching_title = conditionally_add_query_item(matching_title, 'is_folder', is_folder)
    matching_title = conditionally_add_query_item(matching_title, 'is_registration', is_registration)

    if len(ignore_nodes) > 0:
        for node_id in ignore_nodes:
            matching_title = matching_title & Q('_id', 'ne', node_id)

    my_projects = []
    my_project_count = 0
    public_projects = []

    if include_contributed == "yes":
        my_projects = Node.find(
            matching_title &
            Q('contributors', 'contains', user._id)  # user is a contributor
        ).limit(max_results)
        my_project_count = my_project_count

    if my_project_count < max_results and include_public == "yes":
        public_projects = Node.find(
            matching_title &
            Q('is_public', 'eq', True)  # is public
        ).limit(max_results - my_project_count)

    results = list(my_projects) + list(public_projects)
    ret = process_project_search_results(results, **kwargs)
    return ret

@must_be_logged_in
def process_project_search_results(results, **kwargs):
    """
    :param results: list of projects from the modular ODM search
    :return: we return the entire search result, which is a list of
    dictionaries. This includes the list of contributors.
    """
    user = kwargs['auth'].user

    ret = []

    for project in results:
        authors = get_node_contributors_abbrev(project=project, auth=kwargs['auth'])
        authors_html = ''
        for author in authors['contributors']:
            a = User.load(author['user_id'])
            authors_html += '%s' % (a.url, a.fullname)
            authors_html += author['separator'] + ' '
        authors_html += ' ' + authors['others_count']

        ret.append({
            'id': project._id,
            'label': project.title,
            'value': project.title,
            'category': 'My Projects' if user in project.contributors else 'Public Projects',
            'authors': authors_html,
        })

    return ret

@collect_auth
def search_contributor(auth):
    user = auth.user if auth else None
    nid = request.args.get('excludeNode')
    exclude = Node.load(nid).contributors if nid else []
    query = bleach.clean(request.args.get('query', ''), tags=[], strip=True)
    page = int(bleach.clean(request.args.get('page', '0'), tags=[], strip=True))
    size = int(bleach.clean(request.args.get('size', '5'), tags=[], strip=True))
    return search.search_contributor(query=query, page=page, size=size,
                                     exclude=exclude, current_user=user)

@handle_search_errors
def search_share():
    tick = time.time()
    results = {}

    count = request.args.get('count') is not None
    raw = request.args.get('raw') is not None
    version = request.args.get('v')
    index = 'share_v{}'.format(version) if version else 'share'

    if request.method == 'POST':
        query = request.get_json()
    elif request.method == 'GET':
        query = build_query(
            request.args.get('q', '*'),
            request.args.get('from'),
            request.args.get('size'),
            sort=request.args.get('sort')
        )

    if count:
        results = search.count_share(query, index=index)
    else:
        results = search.search_share(query, raw, index=index)

    results['time'] = round(time.time() - tick, 2)
    return results

@handle_search_errors
def search_share_stats():
    q = request.args.get('q')
    query = build_query(q, 0, 0) if q else {}

    return search.share_stats(query=query)

@handle_search_errors
def search_share_atom(**kwargs):
    q = request.args.get('q', '*')
    sort = request.args.get('sort', 'dateUpdated')

    # we want the results per page to be constant between pages
    # TODO -  move this functionality into build_query in util
    start = util.compute_start(request.args.get('page', 1), RESULTS_PER_PAGE)

    query = build_query(q, size=RESULTS_PER_PAGE, start=start, sort=sort)

    try:
        search_results = search.search_share(query, index='share_v1')
    except MalformedQueryError:
        raise HTTPError(http.BAD_REQUEST)
    except IndexNotFoundError:
        search_results = {
            'count': 0,
            'results': []
        }

    atom_url = api_url_for('search_share_atom', _xml=True, _absolute=True)

    return util.create_atom_feed(
        name='SHARE',
        data=search_results['results'],
        query=q,
        size=RESULTS_PER_PAGE,
        start=start,
        url=atom_url,
        to_atom=share_search.to_atom
    )

def search_share_providers():
    return search.share_providers()
<|endoftext|>


<commit_before>""" Main module provides crawling functions and user interface """

from optparse import OptionParser
from logging import getLogger, StreamHandler

#from .attacks import drive_all
from .crawler import Crawler
from .client import Client
from .utils import get_url_host
from .attacks import drive_all, AttackList

EXIT_CODE = 0

def capture_warning():
    """ When called, sets EXIT_CODE to 1"""
    global EXIT_CODE
    EXIT_CODE = 1

class LogHandler(StreamHandler):
    def __init__(self, stream=None):
        StreamHandler.__init__(self)
        # Using Monkeypatching to use old function
        self.real_emit = self.emit
        self.emit = self.handle_emit

    def handle_emit(self, record):
        msg = self.format(record)
        if "Vulnerability" in msg:
            capture_warning()

        self.real_emit(record)

def exit_main():
    """ Returns exit_code, is used for logging. """
    exit(EXIT_CODE)

log = getLogger(__name__)

def main():
    """ The main function. """
    parser = OptionParser(usage='usage: %prog [options] http(s)://target/ '
                                '[http(s)://another.target/]')
    parser.add_option('--no-crawl', action='store_true', dest='no_crawl',
                      help="DO NOT search for links on the target")
    parser.add_option('--whitelist', default=set(), dest="white_list",
                      help="Hosts which are allowed to be crawled.")
    parser.add_option('--auth',  default=None, dest="auth",
                      help="Post target for authentification")
    parser.add_option('--auth-data',  dest='auth_data',
                      action='append', type='str')
    parser.add_option('--verbose', default=None, dest="verbose",
                      action="store_true")

    # Options for scanning for specific vulnerabilities.
    for attack in AttackList():
        parser.add_option('--' + attack.name, dest=attack.name,
                          action="store_true", default=False)

    options, arguments = parser.parse_args()

    attacks = []

    for attack in AttackList():
        if attack.name in options.__dict__:
            if options.__dict__[attack.name]:
                attacks.extend([attack])

    if len(attacks) == 0:
        attacks = AttackList()

    log.addHandler(LogHandler())

    if len(arguments) < 1:
        parser.error('Invalid amount of arguments')

    client = Client()

    if options.auth is not None:
        if options.auth_data is not None:
            post_data = {}

            for field in options.auth_data:
                name, _, value = field.partition('=')
                post_data.update({name: value})

            client.download(options.auth, post_data, False)

    for target in arguments:
        host = get_url_host(target)
        if host not in options.white_list:
            options.white_list.update({host})

        if options.no_crawl:
            page = client.download_page(target)
            if page is not None:
                drive_all(page, attacks)
        else:
            for link in Crawler(target, options.white_list, client):
                if options.verbose:
                    print("Scanning " + link.url)
                drive_all(link, attacks)

    exit_main()
<commit_msg>__init__.py: Added -v option<commit_after>""" Main module provides crawling functions and user interface """

from optparse import OptionParser
from logging import getLogger, StreamHandler

#from .attacks import drive_all
from .crawler import Crawler
from .client import Client
from .utils import get_url_host
from .attacks import drive_all, AttackList

EXIT_CODE = 0

def capture_warning():
    """ When called, sets EXIT_CODE to 1"""
    global EXIT_CODE
    EXIT_CODE = 1

class LogHandler(StreamHandler):
    def __init__(self, stream=None):
        StreamHandler.__init__(self)
        # Using Monkeypatching to use old function
        self.real_emit = self.emit
        self.emit = self.handle_emit

    def handle_emit(self, record):
        msg = self.format(record)
        if "Vulnerability" in msg:
            capture_warning()

        self.real_emit(record)

def exit_main():
    """ Returns exit_code, is used for logging. """
    exit(EXIT_CODE)

log = getLogger(__name__)

def main():
    """ The main function. """
    parser = OptionParser(usage='usage: %prog [options] http(s)://target/ '
                                '[http(s)://another.target/]')
    parser.add_option('--no-crawl', action='store_true', dest='no_crawl',
                      help="DO NOT search for links on the target")
    parser.add_option('--whitelist', default=set(), dest="white_list",
                      help="Hosts which are allowed to be crawled.")
    parser.add_option('--auth',  default=None, dest="auth",
                      help="Post target for authentification")
    parser.add_option('--auth-data',  dest='auth_data',
                      action='append', type='str')
    parser.add_option('--verbose', '-v', default=None, dest="verbose",
                      action="store_true")

    # Options for scanning for specific vulnerabilities.
    for attack in AttackList():
        parser.add_option('--' + attack.name, dest=attack.name,
                          action="store_true", default=False)

    options, arguments = parser.parse_args()

    attacks = []

    for attack in AttackList():
        if attack.name in options.__dict__:
            if options.__dict__[attack.name]:
                attacks.extend([attack])

    if len(attacks) == 0:
        attacks = AttackList()

    log.addHandler(LogHandler())

    if len(arguments) < 1:
        parser.error('Invalid amount of arguments')

    client = Client()

    if options.auth is not None:
        if options.auth_data is not None:
            post_data = {}

            for field in options.auth_data:
                name, _, value = field.partition('=')
                post_data.update({name: value})

            client.download(options.auth, post_data, False)

    for target in arguments:
        host = get_url_host(target)
        if host not in options.white_list:
            options.white_list.update({host})

        if options.no_crawl:
            page = client.download_page(target)
            if page is not None:
                drive_all(page, attacks)
        else:
            for link in Crawler(target, options.white_list, client):
                if options.verbose:
                    print("Scanning " + link.url)
                drive_all(link, attacks)

    exit_main()
<|endoftext|>


<commit_before>import multiprocessing as mulproc

from multiprocessing import Process

class GPUMapperFutures(object):

    def __init__(self, gpu_idxs):
        self.worker_idxs = gpu_idxs

class GPUMapperMultiprocessing(object):

    def __init__(self, gpu_idxs):
        self.worker_idxs = gpu_idxs
        self.free_workers =

    @property
    def n_workers(self):
        return len(self.gpu_idxs)

    def get_worker_idx(self):
        pass

    def release_worker(self, worker_idx):
        pass

    def exec_call(self, func, walker_idx, *args):
        raise NotImplementedError

    def map(self, func, *iterables):

        # structure for the results to go into
        results = mulproc.Manager().list()

        # make the pool of work to be done
        walkers_pool = []
        for walker_idx, args in enumerate(zip(*iterables)):
            walkers_pool.append(Process(target=self.exec_call,
                                        args=(func, walker_idx, *args)))

        #

class GPUMapper(object):

    def __init__(self, n_walkers, n_workers=None, gpu_indices=None):

        if gpu_indices is not None:
            self.gpu_indices = gpu_indices
            self.n_workers = len(gpu_indices)
        else:
            assert n_workers, "If gpu_indices are not given the n_workers must be given"
            self.n_workers = n_workers
            self.gpu_indices = range(n_workers)

        # make a Queue for free workers, when one is being used it is
        # popped off and locked
        self.free_workers = mulproc.Queue()

        # the semaphore provides the locks on the workers
        self.lock = mulproc.Semaphore(self.n_workers)

        # initialize a list to put results in
        self.results_list = mulproc.Manager().list()
        for i in range(n_walkers):
            self.results_list.append(None)

        # add the free worker indices (not device/gpu indices) to the
        # free workers queue
        for i in range(self.n_workers):
            self.free_workers.put(i)

    def exec_call(self, _call_, index, *args):

        # lock a worker, then pop it off the queue so no other process
        # tries to use it
        self.lock.acquire()
        worker_idx = self.free_workers.get()

        # convert this to an available GPU index
        gpu_idx = self.gpu_indices[worker_idx]

        # call the function which is accepted as kwargs
        result = _call_(*args, DeviceIndex=str(gpu_idx))

        self.free_workers.put(worker_idx)
        self.results_list[index] = result
        self.lock.release()

    def map(self, _call_, *iterables):

        walkers_pool = []
        # create processes and start to run
        index = 0

        for args in zip(*iterables):
            walkers_pool.append(mulproc.Process(target=self.exec_call,
                                                args=(_call_, index, *args)))
            index += 1

        for p in walkers_pool:
            p.start()

        for p in walkers_pool:
            p.join()

        return self.results_list
<commit_msg>update<commit_after>import multiprocessing as mulproc

from multiprocessing import Process

class GPUProcess(Process):
    def __init__(self, task_queue, result_queue):
        Process.__init__(self)

        self.task_queue = task_queue
        self.result_queue = result_queue

    def run(self):
        while True:

            # get the next task
            task_idx, next_task = self.task_queue.get()

            # check for the poison pill which is the signal to stop
            if next_task is None:

                print('{}: Exiting'.format(self.name))

                # mark the task as done
                self.task_queue.task_done()

                # and exit the loop
                break

            print('{}: {}'.format(self.name, next_task))

            # run the task
            answer = next_task()

            # (for joinable queue) tell the queue that the formerly
            # enqued task is complete
            self.task_queue.task_done()

            # put the results into the results queue with it's task
            # index so we can sort them later
            self.result_queue.put((task_idx, answer))

class GPUQueueMapper(object):

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

        # initialize the func to None, this must be set a
        self.func = None

    def map(self, tasks):

class GPUMapper(object):

    def __init__(self, n_walkers, n_workers=None, gpu_indices=None):

        if gpu_indices is not None:
            self.gpu_indices = gpu_indices
            self.n_workers = len(gpu_indices)
        else:
            assert n_workers, "If gpu_indices are not given the n_workers must be given"
            self.n_workers = n_workers
            self.gpu_indices = range(n_workers)

        # make a Queue for free workers, when one is being used it is
        # popped off and locked
        self.free_workers = mulproc.Queue()

        # the semaphore provides the locks on the workers
        self.lock = mulproc.Semaphore(self.n_workers)

        # initialize a list to put results in
        self.results_list = mulproc.Manager().list()
        for i in range(n_walkers):
            self.results_list.append(None)

        # add the free worker indices (not device/gpu indices) to the
        # free workers queue
        for i in range(self.n_workers):
            self.free_workers.put(i)

    def exec_call(self, _call_, index, *args):

        # lock a worker, then pop it off the queue so no other process
        # tries to use it
        self.lock.acquire()
        worker_idx = self.free_workers.get()

        # convert this to an available GPU index
        gpu_idx = self.gpu_indices[worker_idx]

        # call the function which is accepted as kwargs
        result = _call_(*args, DeviceIndex=str(gpu_idx))

        self.free_workers.put(worker_idx)
        self.results_list[index] = result
        self.lock.release()

    def map(self, _call_, *iterables):

        walkers_pool = []
        # create processes and start to run
        index = 0

        for args in zip(*iterables):
            walkers_pool.append(mulproc.Process(target=self.exec_call,
                                                args=(_call_, index, *args)))
            index += 1

        for p in walkers_pool:
            p.start()

        for p in walkers_pool:
            p.join()

        return self.results_list
<|endoftext|>


<commit_before>from django.db import models
from django import forms

from wheelcms_axle.content import type_registry
from wheelcms_axle.content import Content
from wheelcms_axle.templates import template_registry
from wheelcms_axle.spoke import Spoke
from wheelcms_axle.forms import formfactory

from tinymce.widgets import TinyMCE
from tinymce.models import HTMLField

##
## Either HTMLField or TinyMCE widget. But a demonstration on how to alter
## a Spoke's form is also nice to have.

class Page(Content):
    """ A simple page object """
    body = models.TextField(blank=False)
    # body = HTMLField(blank=False)

class PageForm(formfactory(Page)):
    body = forms.CharField(widget=TinyMCE(), required=False)

class PageType(Spoke):
    document_fields = Spoke.document_fields + ("body", )

    model = Page
    title = "A simple HTML page"
    form = PageForm

    @property
    def icon(self):
        if self.instance.node.children():
            return "folder_page.png"
        return "page.png"

type_registry.register(PageType)
template_registry.register(PageType, "wheelcms_spokes/page_view.html",
                           "Basic Page view", default=True)
<commit_msg>possibly more efficient way to check children<commit_after>from django.db import models
from django import forms

from wheelcms_axle.content import type_registry
from wheelcms_axle.content import Content
from wheelcms_axle.templates import template_registry
from wheelcms_axle.spoke import Spoke
from wheelcms_axle.forms import formfactory

from tinymce.widgets import TinyMCE
from tinymce.models import HTMLField

##
## Either HTMLField or TinyMCE widget. But a demonstration on how to alter
## a Spoke's form is also nice to have.

class Page(Content):
    """ A simple page object """
    body = models.TextField(blank=False)
    # body = HTMLField(blank=False)

class PageForm(formfactory(Page)):
    body = forms.CharField(widget=TinyMCE(), required=False)

class PageType(Spoke):
    document_fields = Spoke.document_fields + ("body", )

    model = Page
    title = "A simple HTML page"
    form = PageForm

    @property
    def icon(self):
        if self.instance.node.children().exists():
            return "folder_page.png"
        return "page.png"

type_registry.register(PageType)
template_registry.register(PageType, "wheelcms_spokes/page_view.html",
                           "Basic Page view", default=True)
<|endoftext|>


<commit_before>VERSION = (0, 7, 6)

def version_str():
    if len(VERSION) == 3:
        return "%s.%s.%s" % VERSION
    elif len(VERSION) == 4:
        return "%s.%s.%s-%s" % VERSION
    else:
        raise IndexError("Incorrect format for the VERSION tuple")<commit_msg>update version to 0.7.8<commit_after>VERSION = (0, 7, 8)

def version_str():
    if len(VERSION) == 3:
        return "%s.%s.%s" % VERSION
    elif len(VERSION) == 4:
        return "%s.%s.%s-%s" % VERSION
    else:
        raise IndexError("Incorrect format for the VERSION tuple")<|endoftext|>


<commit_before>VERSION = (0, 6, 0)

def version_str():
    if len(VERSION) == 3:
        return "%s.%s.%s" % VERSION
    elif len(VERSION) == 4:
        return "%s.%s.%s-%s" % VERSION
    else:
        raise IndexError("Incorrect format for the VERSION tuple")<commit_msg>update version to 0.6.1<commit_after>VERSION = (0, 6, 1)

def version_str():
    if len(VERSION) == 3:
        return "%s.%s.%s" % VERSION
    elif len(VERSION) == 4:
        return "%s.%s.%s-%s" % VERSION
    else:
        raise IndexError("Incorrect format for the VERSION tuple")<|endoftext|>


<commit_before>from __future__ import division
from trie import Trie
from sys import argv
from time import clock
from random import choice

trie = Trie()
arr = []
n = 1000
if len(argv) > 1:
	n = int(argv[1])

words = [line for line in open("benchmarking/wordlist")][:n]

def insert_trie():
	for i in xrange(n):
		trie.insert(words[i])

def insert_array():
	for i in xrange(n):
		arr.append(words[i])

def trie_autocomplete():
	word = choice(words)[:3]
	trie.autocomplete(word)

def array_autocomplete():
	for i in xrange(n):
		word = choice(words)[:3]
		trie.autocomplete(word)

def timefn(f, times):
	for i in xrange(n):
		total = 0
		for i in xrange(times):
			start = clock()
			f()
			end = clock()
			total += end - start
		return total / times

def benchmark():
	print "array insertion: %.10f" % timefn(insert_array, 5)
	print "trie insertion: %.10f" % timefn(insert_trie, 5)
	print "array autocomplete: %.10f" % timefn(array_autocomplete, 5)
	print "trie autocomplete: %.10f" % timefn(trie_autocomplete, 5)

if __name__ == '__main__':
	benchmark()
<commit_msg>python benchmark<commit_after>from __future__ import division
from trie import Trie
from sys import argv
from time import clock
from random import choice

trie = Trie()
arr = []
n = 1000
if len(argv) > 1:
	n = int(argv[1])

words = [line for line in open("benchmarking/wordlist")][:n]

def insert_trie():
	for i in xrange(n):
		trie.insert(words[i])

def insert_array():
	for i in xrange(n):
		arr.append(words[i])

def trie_autocomplete():
	for i in xrange(n):
		word = choice(words)[:3]
		trie.autocomplete(word)

def array_autocomplete():
	for i in xrange(n):
		word = choice(words)[:3]
		a = [w for w in words if w.startswith(word)]

def timefn(f, times):
	for i in xrange(n):
		total = 0
		for i in xrange(times):
			start = clock()
			f()
			end = clock()
			total += end - start
		return total / times

def benchmark():
	print "array insertion: %.10f" % timefn(insert_array, 5)
	print "trie insertion: %.10f" % timefn(insert_trie, 5)
	print "array autocomplete: %.10f" % timefn(array_autocomplete, 1)
	print "trie autocomplete: %.10f" % timefn(trie_autocomplete, 1)

if __name__ == '__main__':
	benchmark()
<|endoftext|>


<commit_before>import json
import os
import copy
import shutil

from django.http import (
    HttpResponse,
    HttpResponseNotAllowed,
    HttpResponseForbidden,
    HttpResponseBadRequest,
    HttpResponseNotFound
)

from wand.image import Image as WandImage

from .conf.app import settings
from .models import Image, source_upload_to

ACC_HEADERS = {
    'Access-Control-Allow-Origin': '*',
    'Access-Control-Allow-Methods': 'GET, POST, OPTIONS',
    'Access-Control-Max-Age': 1000,
    'Access-Control-Allow-Headers': '*'
}

def crossdomain(origin="*", methods=[], headers=["X-Betty-Api-Key", "Content-Type"]):

    def _method_wrapper(func):

        def _crossdomain_wrapper(request, *args, **kwargs):
            if request.method != "OPTIONS":
                response = func(request, *args, **kwargs)
            else:
                response = HttpResponse()
            response["Access-Control-Allow-Origin"] = "*"
            if methods:
                if request.method not in methods:
                    return HttpResponseNotAllowed(methods)
                response["Access-Control-Allow-Methods"] = ", ".join(methods)
            if headers:
                response["Access-Control-Allow-Headers"] = ", ".join(headers)
            return response

        return _crossdomain_wrapper

    return _method_wrapper

@crossdomain(methods=['POST', 'OPTIONS'])
def new(request):
    if not (request.user.is_superuser or request.user.has_perm("betty.add_image")):
        response_text = json.dumps({'message': 'Not authorized'})
        return HttpResponseForbidden(response_text, content_type="application/json")

    image_file = request.FILES.get("image")
    if image_file is None:
        return HttpResponseBadRequest()

    image = Image.objects.create(name=image_file.name)
    os.makedirs(image.path())
    source_path = source_upload_to(image, image_file.name)

    with open(source_path, 'wb+') as f:
        for chunk in image_file.chunks():
            f.write(chunk)
        f.seek(0)
        with WandImage(file=f) as img:
            image.width = img.size[0]
            image.height = img.size[1]

        image.source.name = source_path
        image.save()

    return HttpResponse(json.dumps(image.to_native()), content_type="application/json")

@crossdomain(methods=['POST', 'OPTIONS'])
def update_selection(request, image_id, ratio_slug):
    if not (request.user.is_superuser or request.user.has_perm("betty.change_image")):
        response_text = json.dumps({'message': 'Not authorized'})
        return HttpResponseForbidden(response_text, content_type="application/json")

    try:
        image = Image.objects.get(id=image_id)
    except Image.DoesNotExist:
        message = json.dumps({"message": "No such image!"})
        return HttpResponseNotFound(message, content_type="application/json")

    try:
        request_json = json.loads(request.body)
    except Exception:
        message = json.dumps({"message": "Bad selection"})
        return HttpResponseBadRequest(message, content_type="application/json")
    try:
        selection = {
            "x0": int(request_json["x0"]),
            "y0": int(request_json["y0"]),
            "x1": int(request_json["x1"]),
            "y1": int(request_json["y1"]),
        }
    except (KeyError, ValueError):
        message = json.dumps({"message": "Bad selection"})
        return HttpResponseBadRequest(message, content_type="application/json")

    selections = copy.copy(image.selections)
    if selections is None:
        selections = {}

    if ratio_slug not in settings.BETTY_RATIOS:
        message = json.dumps({"message": "No such ratio"})
        return HttpResponseBadRequest(message, content_type="application/json")

    selections[ratio_slug] = selection
    image.selections = selections
    image.save()

    # TODO: Use a celery task for this?
    ratio_path = os.path.join(image.path(), ratio_slug)
    if os.path.exists(ratio_path):
        # crops = os.listdir(ratio_path)
        # TODO: flush cache on crops
        shutil.rmtree(ratio_path)

    message = json.dumps({"message": "Update sucessfull"})
    return HttpResponse(message, content_type="application/json")

@crossdomain(methods=['GET', 'OPTIONS'])
def search(request):
    if not request.user.is_staff:
        response_text = json.dumps({'message': 'Not authorized'})
        return HttpResponseForbidden(response_text, content_type="application/json")

    results = []
    query = request.GET.get("q")
    if query:
        for image in Image.objects.filter(name__icontains=query)[:20]:
            results.append(image.to_native())
    return HttpResponse(json.dumps(results), content_type="application/json")

@crossdomain(methods=["GET", "PATCH", "OPTIONS"])
def detail(request, image_id):
    if not (request.user.is_superuser or request.user.has_perm("betty.change_image")):
        response_text = json.dumps({'message': 'Not authorized'})
        return HttpResponseForbidden(response_text, content_type="application/json")

    try:
        image = Image.objects.get(id=image_id)
    except Image.DoesNotExist:
        message = json.dumps({"message": "No such image!"})
        return HttpResponseNotFound(message, content_type="application/json")

    if request.method == "PATCH":
        try:
            request_json = json.loads(request.body)
        except Exception:
            message = json.dumps({"message": "Bad Request"})
            return HttpResponseBadRequest(message, content_type="application/json")

        update_fields = []
        for field in ("name", "credit", "selections"):
            if field in request_json:
                setattr(image, field, request_json[field])
                update_fields.append(field)
        image.save(update_fields=update_fields)

    return HttpResponse(json.dumps(image.to_native()), content_type="application/json")
<commit_msg>Simlifying auth checking<commit_after>import json
import os
import copy
import shutil

from django.http import (
    HttpResponse,
    HttpResponseNotAllowed,
    HttpResponseForbidden,
    HttpResponseBadRequest,
    HttpResponseNotFound
)

from wand.image import Image as WandImage

from .conf.app import settings
from .models import Image, source_upload_to

ACC_HEADERS = {
    'Access-Control-Allow-Origin': '*',
    'Access-Control-Allow-Methods': 'GET, POST, OPTIONS',
    'Access-Control-Max-Age': 1000,
    'Access-Control-Allow-Headers': '*'
}

def crossdomain(origin="*", methods=[], headers=["X-Betty-Api-Key", "Content-Type"]):

    def _method_wrapper(func):

        def _crossdomain_wrapper(request, *args, **kwargs):
            if request.method != "OPTIONS":
                response = func(request, *args, **kwargs)
            else:
                response = HttpResponse()
            response["Access-Control-Allow-Origin"] = "*"
            if methods:
                if request.method not in methods:
                    return HttpResponseNotAllowed(methods)
                response["Access-Control-Allow-Methods"] = ", ".join(methods)
            if headers:
                response["Access-Control-Allow-Headers"] = ", ".join(headers)
            return response

        return _crossdomain_wrapper

    return _method_wrapper

@crossdomain(methods=['POST', 'OPTIONS'])
def new(request):
    if not request.user.has_perm("betty.add_image"):
        response_text = json.dumps({'message': 'Not authorized'})
        return HttpResponseForbidden(response_text, content_type="application/json")

    image_file = request.FILES.get("image")
    if image_file is None:
        return HttpResponseBadRequest()

    image = Image.objects.create(name=image_file.name)
    os.makedirs(image.path())
    source_path = source_upload_to(image, image_file.name)

    with open(source_path, 'wb+') as f:
        for chunk in image_file.chunks():
            f.write(chunk)
        f.seek(0)
        with WandImage(file=f) as img:
            image.width = img.size[0]
            image.height = img.size[1]

        image.source.name = source_path
        image.save()

    return HttpResponse(json.dumps(image.to_native()), content_type="application/json")

@crossdomain(methods=['POST', 'OPTIONS'])
def update_selection(request, image_id, ratio_slug):
    if not request.user.has_perm("betty.change_image"):
        response_text = json.dumps({'message': 'Not authorized'})
        return HttpResponseForbidden(response_text, content_type="application/json")

    try:
        image = Image.objects.get(id=image_id)
    except Image.DoesNotExist:
        message = json.dumps({"message": "No such image!"})
        return HttpResponseNotFound(message, content_type="application/json")

    try:
        request_json = json.loads(request.body)
    except Exception:
        message = json.dumps({"message": "Bad selection"})
        return HttpResponseBadRequest(message, content_type="application/json")
    try:
        selection = {
            "x0": int(request_json["x0"]),
            "y0": int(request_json["y0"]),
            "x1": int(request_json["x1"]),
            "y1": int(request_json["y1"]),
        }
    except (KeyError, ValueError):
        message = json.dumps({"message": "Bad selection"})
        return HttpResponseBadRequest(message, content_type="application/json")

    selections = copy.copy(image.selections)
    if selections is None:
        selections = {}

    if ratio_slug not in settings.BETTY_RATIOS:
        message = json.dumps({"message": "No such ratio"})
        return HttpResponseBadRequest(message, content_type="application/json")

    selections[ratio_slug] = selection
    image.selections = selections
    image.save()

    # TODO: Use a celery task for this?
    ratio_path = os.path.join(image.path(), ratio_slug)
    if os.path.exists(ratio_path):
        # crops = os.listdir(ratio_path)
        # TODO: flush cache on crops
        shutil.rmtree(ratio_path)

    message = json.dumps({"message": "Update sucessfull"})
    return HttpResponse(message, content_type="application/json")

@crossdomain(methods=['GET', 'OPTIONS'])
def search(request):
    if not request.user.is_staff:
        response_text = json.dumps({'message': 'Not authorized'})
        return HttpResponseForbidden(response_text, content_type="application/json")

    results = []
    query = request.GET.get("q")
    if query:
        for image in Image.objects.filter(name__icontains=query)[:20]:
            results.append(image.to_native())
    return HttpResponse(json.dumps(results), content_type="application/json")

@crossdomain(methods=["GET", "PATCH", "OPTIONS"])
def detail(request, image_id):
    if not request.user.has_perm("betty.change_image"):
        response_text = json.dumps({'message': 'Not authorized'})
        return HttpResponseForbidden(response_text, content_type="application/json")

    try:
        image = Image.objects.get(id=image_id)
    except Image.DoesNotExist:
        message = json.dumps({"message": "No such image!"})
        return HttpResponseNotFound(message, content_type="application/json")

    if request.method == "PATCH":
        try:
            request_json = json.loads(request.body)
        except Exception:
            message = json.dumps({"message": "Bad Request"})
            return HttpResponseBadRequest(message, content_type="application/json")

        update_fields = []
        for field in ("name", "credit", "selections"):
            if field in request_json:
                setattr(image, field, request_json[field])
                update_fields.append(field)
        image.save(update_fields=update_fields)

    return HttpResponse(json.dumps(image.to_native()), content_type="application/json")
<|endoftext|>