<commit_before>#!/usr/bin/env python
from __future__ import print_function

# native Python imports
import os.path
import time
import sys
import sqlite3
import gzip
import itertools as it

import numpy as np
import toml # toml.py

# third-party imports
import cyvcf as vcf

import zlib
import cPickle

def pack_blob(obj):
    return buffer(zlib.compress(cPickle.dumps(obj, cPickle.HIGHEST_PROTOCOL)))

def is_number(op, field):
    return field.endswith("_float") or op in ("mean", "median", "min", "max")

def is_integer(op, field):
    return field.endswith("_int")

def is_flag(op, infos):
    return infos[op].type == "Flag"

def name_type(op, field):
    """
    >>> name_type("sum", "t_float")
    ('t', 'REAL')
    >>> name_type("aaa", "tt_int")
    ('tt', 'INTEGER')
    >>> name_type("aaa", "ingt_flag")
    ('ingt', 'BOOL')
    >>> name_type("aaa", "hello")
    ('hello', 'TEXT')
    """
    if field.endswith("_float"):
        return field[:-6], "REAL"
    if field.endswith("_int"):
        return field[:-4], "INTEGER"
    if field.endswith("flag"):
        return field[:-5], "BOOL"
    if op == "flag":
        return field, "BOOL"
    if op in ("min", "mean", "max"):
        return field, "REAL"

    return field, "TEXT"

def xopen(path, mode="r"):
    if path == "-":
        if mode[0] == "w":
            return sys.stdout
        return sys.stdin

    if path.endswith(".gz"):
        return gzip.open(path, mode)
    return open(path, mode)

class VCFLoader(object):
    def __init__(self, path, toml_path, db_path="test.db"):

        # this will find the fields we are expecting during load.
        self.load_config(toml_path)
        self.vcf_reader = vcf.Reader(open(path))
        if db_path is None:
            db_path = (path[:-6 if path.endswith('.vcf.gz') else -3]) + "db"
        self.db_path = db_path
        self.db = sqlite3.connect(self.db_path)
        self.db.isolation_level = None
        self.cursor = self.db.cursor()
        self.cursor.execute('PRAGMA synchronous = OFF')
        self.cursor.execute('PRAGMA journal_mode=MEMORY')

        self.create_tables()

        self.extended_fields = ["chrom", "start", "end", "variant_id", "vid", "ref",
                                "alt", "qual", "filter", "type", "sub_type"] \
                               + self.fields + ["gts", "gt_types", "gt_phases",
                                       "gt_depths", "gt_ref_depths",
                                       "gt_alt_depths", "gt_quals",
                                       "gt_copy_numbers", "gt_phred_ll_homref",
                                       "gt_phred_ll_het", "gt_phred_ll_homalt"]
    def create_tables(self):
        self._create_variants_table()
        self._create_sample_table()

    def _create_sample_table(self):
        samples = self.vcf_reader.samples
        self.sample_to_idx = {s: i for i, s in enumerate(samples, start=1)}
        self.cursor.execute("""\
CREATE TABLE samples (
    sample_id INTEGER,
    family_id text default NULL,
    name text default NULL,
    paternal_id text default NULL,
    maternal_id TEXT default NULL,
    sex TEXT default NULL,
    phenotype TEXT default NULL
    )""")
        q = "INSERT INTO samples VALUES (%s)" % ",".join(["?"] * 7)
        self.cursor.executemany(q, [(self.sample_to_idx[s], 0, s, 0, 0, -9, -9) for s in samples])

    def _create_variants_table(self):
        tmpl = """\
CREATE table variants (
    chrom TEXT,
    start INTEGER,
    end INTEGER,
    variant_id INTEGER,
    vid TEXT,
    ref TEXT,
    alt TEXT,
    qual REAL,
    filter TEXT,
    type TEXT,
    sub_type TEXT,
    %s,
    gts BLOB,
    gt_types BLOB,
    gt_phases BLOB,
    gt_depths BLOB,
    gt_ref_depths BLOB,
    gt_alt_depths BLOB,
    gt_quals BLOB,
    gt_copy_numbers BLOB,
    gt_phred_ll_homref BLOB,
    gt_phred_ll_het BLOB,
    gt_phred_ll_homalt BLOB,
    PRIMARY KEY(variant_id ASC)
    )"""
        cmd = tmpl % (",\n    ".join("%s %s" % (field, ftype) for field, ftype in
                      it.izip(self.fields, self.types)))
        self.cursor.execute(cmd)

    def load_config(self, toml_path):
        cfg = toml.loads(xopen(toml_path).read())
        fields = []
        ops = []
        for f in cfg['annotation']:
            fields.extend(f.get('names', f.get('fields')))
            ops.extend(f['ops'])
        assert len(ops) == len(fields)
        self.config = cfg
        self.o_fields = fields
        self.ops = ops

        field_types = [name_type(o, f) for f, o in it.izip(self.o_fields, self.ops)]

        self.fields, self.types = map(list, zip(*field_types))

    def insert(self, variants):
        if len(variants) == 0:
            return
        t0 = time.time()
        cur = self.cursor
        query = "INSERT INTO variants values(%s)" % ",".join(["?"] * len(variants[0]))
        try:
            cur.executemany(query, variants)
        except sqlite3.ProgrammingError:
            raise
            for var in variants:
                cur.execute(query, var)
        print("insert time:", time.time() - t0)

    def load(self, buffer_size=10000):
        tl = time.time()
        load_buffer = []
        ots = zip(self.types, self.o_fields)
        tgt, igt,bt = 0, 0, 0
        for i, v in enumerate(self.vcf_reader, start=1):
            t0 = time.time()
            vals = [v.INFO.get(f, '' if t == 'TEXT' else False if t == 'BOOL' else  None) for t, f in ots]
            igt += time.time() - t0

            t0 = time.time()
            gts = [
                pack_blob(np.array(v.gt_bases, np.str)),
                pack_blob(np.array(v.gt_types, np.int8)),
                pack_blob(np.array(v.gt_phases, np.bool)),
                pack_blob(np.array(v.gt_depths, np.int32)),
                pack_blob(np.array(v.gt_ref_depths, np.int32)),
                pack_blob(np.array(v.gt_alt_depths, np.int32)),
                pack_blob(np.array(v.gt_quals, np.float32)),
                pack_blob(np.array(v.gt_copy_numbers, np.float32)),
                pack_blob(None),
                pack_blob(None),
                pack_blob(None), ]
            tgt += time.time() - t0

            t0 = time.time()
            basic = [v.CHROM, v.start, v.end, i - 1, v.ID or '', v.REF,
                     ",".join(v.ALT), v.QUAL, v.FILTER, v.var_type,
                     v.var_subtype]
            bt += time.time() - t0
            variant = basic + vals + gts
            load_buffer.append(variant)

            if i % buffer_size == 0:
                self.insert(load_buffer)
                load_buffer = []
        self.insert(load_buffer)
        print("gt time", tgt)
        print("info time", igt)
        print("basic time", bt)
        print("total load time:", time.time() - tl)

if __name__ == "__main__":
    import doctest
    doctest.testmod()

    v = VCFLoader(sys.argv[1], sys.argv[2])
    v.load()
<commit_msg>add more fields to the db with the new loader.<commit_after>#!/usr/bin/env python
from __future__ import print_function
from effects import SnpEff

# native Python imports
import os.path
import time
import sys
import sqlite3
import gzip
import itertools as it

import numpy as np
import toml # toml.py

# third-party imports
import cyvcf as vcf

import zlib
import cPickle

def pack_blob(obj, _none=buffer(zlib.compress(cPickle.dumps(None, cPickle.HIGHEST_PROTOCOL)))):
    if obj is None: return _none
    return buffer(zlib.compress(cPickle.dumps(obj, cPickle.HIGHEST_PROTOCOL)))

def is_number(op, field):
    return field.endswith("_float") or op in ("mean", "median", "min", "max")

def is_integer(op, field):
    return field.endswith("_int")

def is_flag(op, infos):
    if not op in infos: return False
    return infos[op].type == "Flag"

def name_type(op, field, infos):
    """
    >>> name_type("sum", "t_float", {})
    ('t', 'REAL')
    >>> name_type("aaa", "tt_int", {})
    ('tt', 'INTEGER')
    >>> name_type("aaa", "ingt_flag", {})
    ('ingt', 'BOOL')
    >>> name_type("aaa", "hello", {})
    ('hello', 'TEXT')
    >>> name_type("flag", "LCR", {})
    ('LCR', 'BOOL')
    """
    if field.endswith("_float"):
        return field[:-6], "REAL"
    if field.endswith("_int"):
        return field[:-4], "INTEGER"
    if field.endswith("flag"):
        return field[:-5], "BOOL"
    if is_flag(field, infos):
        return field, "BOOL"
    if op == "flag":
        return field, "BOOL"
    if op in ("min", "mean", "max"):
        return field, "REAL"

    return field, "TEXT"

def xopen(path, mode="r"):
    if path == "-":
        if mode[0] == "w":
            return sys.stdout
        return sys.stdin

    if path.endswith(".gz"):
        return gzip.open(path, mode)
    return open(path, mode)

class VCFLoader(object):
    def __init__(self, path, toml_path, db_path="test.db"):

        # this will find the fields we are expecting during load.
        self.vcf_reader = vcf.Reader(open(path))
        self.load_config(toml_path)
        if db_path is None:
            db_path = (path[:-6 if path.endswith('.vcf.gz') else -3]) + "db"
        self.db_path = db_path
        self.db = sqlite3.connect(self.db_path)
        self.db.isolation_level = None
        self.cursor = self.db.cursor()
        self.cursor.execute('PRAGMA synchronous = OFF')
        self.cursor.execute('PRAGMA journal_mode=MEMORY')

        self.create_tables()

        self.extended_fields = ["chrom", "start", "end", "variant_id", "vid", "ref",
                                "alt", "qual", "filter", "type", "sub_type"] \
                               + self.fields + ["gts", "gt_types", "gt_phases",
                                       "gt_depths", "gt_ref_depths",
                                       "gt_alt_depths", "gt_quals",
                                       #"gt_copy_numbers",
                                       "gt_phred_ll_homref", "gt_phred_ll_het", "gt_phred_ll_homalt"]
    def create_tables(self):
        self._create_variants_table()
        self._create_sample_table()

    def _create_sample_table(self):
        samples = self.vcf_reader.samples
        self.sample_to_idx = {s: i for i, s in enumerate(samples, start=1)}
        self.cursor.execute("""\
CREATE TABLE samples (
    sample_id INTEGER,
    family_id text default NULL,
    name text default NULL,
    paternal_id text default NULL,
    maternal_id TEXT default NULL,
    sex TEXT default NULL,
    phenotype TEXT default NULL
    )""")
        q = "INSERT INTO samples VALUES (%s)" % ",".join(["?"] * 7)
        self.cursor.executemany(q, [(self.sample_to_idx[s], 0, s, 0, 0, -9, -9) for s in samples])

    def _create_variants_table(self):
        tmpl = """\
CREATE table variants (
    chrom TEXT,
    start INTEGER,
    end INTEGER,
    variant_id INTEGER,
    vid TEXT,
    ref TEXT,
    alt TEXT,
    qual REAL,
    filter TEXT,
    type TEXT,
    sub_type TEXT,
    gene TEXT,
    impact_severity TEXT,
    is_exonic BOOL,
    is_coding BOOL,
    is_lof BOOL,
    biotype TEXT,
    sift_pred TEXT,
    sift_score REAL,
    polyphen_pred TEXT,
    polyphen_score REAL,
    %s,
    gts BLOB,
    gt_types BLOB,
    gt_phases BLOB,
    gt_depths BLOB,
    gt_ref_depths BLOB,
    gt_alt_depths BLOB,
    gt_quals BLOB,
    --gt_copy_numbers BLOB,
    gt_phred_ll_homref BLOB,
    gt_phred_ll_het BLOB,
    gt_phred_ll_homalt BLOB,
    PRIMARY KEY(variant_id ASC)
    )"""
        cmd = tmpl % (",\n    ".join("%s %s" % (field, ftype) for field, ftype in
                      it.izip(self.fields, self.types)))
        self.cursor.execute(cmd)

    def load_config(self, toml_path):
        cfg = toml.loads(xopen(toml_path).read())
        fields = []
        ops = []
        for f in cfg['annotation']:
            fields.extend(f.get('names', f.get('fields')))
            ops.extend(f['ops'])
        assert len(ops) == len(fields)
        self.config = cfg
        self.o_fields = fields
        self.ops = ops

        field_types = [name_type(o, f, self.vcf_reader.infos) for f, o in it.izip(self.o_fields, self.ops)]
        self.fields, self.types = map(list, zip(*field_types))

    def insert(self, variants):
        if len(variants) == 0:
            return
        t0 = time.time()
        cur = self.cursor
        query = "INSERT INTO variants values(%s)" % ",".join(["?"] * len(variants[0]))
        try:
            cur.executemany(query, variants)
        except sqlite3.ProgrammingError:
            raise
            for var in variants:
                cur.execute(query, var)
        print("insert time:", time.time() - t0)

    def load(self, buffer_size=10000):
        tl = time.time()
        load_buffer = []
        ots = zip(self.types, self.o_fields)
        tgt, igt, bt = 0, 0, 0
        for i, v in enumerate(self.vcf_reader, start=1):
            t0 = time.time()
            vals = [v.INFO.get(f, '' if t == 'TEXT' else False if t == 'BOOL' else  None) for t, f in ots]
            igt += time.time() - t0

            t0 = time.time()
            gts = [
                pack_blob(np.array(v.gt_bases, np.str)),
                pack_blob(np.array(v.gt_types, np.int8)),
                pack_blob(np.array(v.gt_phases, np.bool)),
                pack_blob(np.array(v.gt_depths, np.int32)),
                pack_blob(np.array(v.gt_ref_depths, np.int32)),
                pack_blob(np.array(v.gt_alt_depths, np.int32)),
                pack_blob(np.array(v.gt_quals, np.float32)),
                #pack_blob(np.array(v.gt_copy_numbers, np.float32)),
                pack_blob(None),
                pack_blob(None),
                pack_blob(None), ]
            tgt += time.time() - t0

            t0 = time.time()
            ann = SnpEff.top_severity(v.INFO['ANN'].split(","))
            if isinstance(ann, list):
                ann = ann[0]
            basic = [v.CHROM, v.start, v.end, i - 1, v.ID or '', v.REF,
                     ",".join(v.ALT), v.QUAL, v.FILTER, v.var_type,
                     v.var_subtype, ann.gene, ann.impact_severity,
                     ann.exonic, ann.coding, ann.lof, ann.biotype,
                     ann.sift_class, ann.sift_value,
                     ann.polyphen_class, ann.polyphen_value]

            bt += time.time() - t0
            variant = basic + vals + gts
            load_buffer.append(variant)

            if i % buffer_size == 0:
                self.insert(load_buffer)
                load_buffer = []
        self.insert(load_buffer)
        print("gt time", tgt)
        print("info time", igt)
        print("basic time", bt)
        print("total load time:", time.time() - tl)

if __name__ == "__main__":
    import doctest
    doctest.testmod()

    v = VCFLoader(sys.argv[1], sys.argv[2])
    v.load()
<|endoftext|>


<commit_before>import datetime
from flask import abort, Blueprint, redirect, render_template, request, url_for
from flask_login import current_user, login_required
from playhouse.flask_utils import get_object_or_404
from werkzeug.datastructures import CombinedMultiDict
from app import app, images
from app.forms import EditEventForm
from app.models import Event

mod = Blueprint('events', __name__, url_prefix='/konserter')

def url_for_other_page(page):
    args = request.view_args.copy()
    args['page'] = page
    return url_for(request.endpoint, **args)

app.jinja_env.globals['url_for_other_page'] = url_for_other_page
app.jinja_env.globals['image_url'] = images.url

@mod.route('/', defaults={'page': 1})
@mod.route('/page/<int:page>/')
def overview(page):
    events = Event.select().order_by(Event.start_time.desc())

    if not current_user.is_authenticated:
        events = events.where(Event.published == True)

    pagination = events.paginate(page, 5)

    if not pagination and events:
        last_page = (len(events) // 5) + 1
        if len(events) % 5:
            last_page += 1
        return redirect(url_for('.konserter', page=last_page))

    has_next = True if events.paginate(page+1, 5) else False

    return render_template('events/overview.html',
                           pagination=pagination,
                           page=page,
                           has_next=has_next)

@mod.route('/new-event/', methods=['GET', 'POST'])
@login_required
def new_event():
    form = EditEventForm(CombinedMultiDict((request.form, request.files)))
    if form.validate_on_submit():
        event = Event.create(
                title=form.title.data,
                path=url_for('.overview'),
                content=form.content.data,
                published=form.published.data,
                start_time=form.start_time.data,
                location=form.location.data,
                timestamp=datetime.datetime.now(),
                author=current_user.id,
                image=images.save(form.upload.data)
                )
        return redirect(url_for('.view_event', slug=event.slug))

    return render_template('events/edit-event.html', form=form)

@mod.route('/<slug>/')
def view_event(slug):
    event = get_object_or_404(Event, Event.slug == slug)

    if not event.published and not current_user.is_authenticated:
        return abort(404)

    return render_template('events/view-event.html', event=event)

@mod.route('/<slug>/edit/', methods=['GET', 'POST'])
@login_required
def edit_event(slug):
    event = get_object_or_404(Event, Event.slug == slug)
    form = EditEventForm(
            CombinedMultiDict((request.form, request.files)),
            event
            )

    if form.validate_on_submit():
        event.title = form.title.data
        event.content = form.content.data
        event.published = form.published.data
        event.start_time = form.start_time.data
        event.location = form.location.data
        event.image = images.save(form.upload.data)
        event.save()
        return redirect(url_for('.view_event', slug=event.slug))

    return render_template('events/edit-event.html', form=form)

@mod.route('/<slug>/remove/')
@login_required
def remove_event(slug):
    event = get_object_or_404(Event, Event.slug == slug)
    event.delete_instance()
    return redirect(url_for('.overview'))
<commit_msg>allow editing events without image upload<commit_after>import datetime
from flask import abort, Blueprint, redirect, render_template, request, url_for
from flask_login import current_user, login_required
from playhouse.flask_utils import get_object_or_404
from werkzeug.datastructures import CombinedMultiDict
from app import app, images
from app.forms import EditEventForm
from app.models import Event

mod = Blueprint('events', __name__, url_prefix='/konserter')

def url_for_other_page(page):
    args = request.view_args.copy()
    args['page'] = page
    return url_for(request.endpoint, **args)

app.jinja_env.globals['url_for_other_page'] = url_for_other_page
app.jinja_env.globals['image_url'] = images.url

@mod.route('/', defaults={'page': 1})
@mod.route('/page/<int:page>/')
def overview(page):
    events = Event.select().order_by(Event.start_time.desc())

    if not current_user.is_authenticated:
        events = events.where(Event.published == True)

    pagination = events.paginate(page, 5)

    if not pagination and events:
        last_page = (len(events) // 5) + 1
        if len(events) % 5:
            last_page += 1
        return redirect(url_for('.konserter', page=last_page))

    has_next = True if events.paginate(page+1, 5) else False

    return render_template('events/overview.html',
                           pagination=pagination,
                           page=page,
                           has_next=has_next)

@mod.route('/new-event/', methods=['GET', 'POST'])
@login_required
def new_event():
    form = EditEventForm(CombinedMultiDict((request.form, request.files)))
    if form.validate_on_submit():
        event = Event.create(
                title=form.title.data,
                path=url_for('.overview'),
                content=form.content.data,
                published=form.published.data,
                start_time=form.start_time.data,
                location=form.location.data,
                timestamp=datetime.datetime.now(),
                author=current_user.id,
                image=images.save(form.upload.data)
                )
        return redirect(url_for('.view_event', slug=event.slug))

    return render_template('events/edit-event.html', form=form)

@mod.route('/<slug>/')
def view_event(slug):
    event = get_object_or_404(Event, Event.slug == slug)

    if not event.published and not current_user.is_authenticated:
        return abort(404)

    return render_template('events/view-event.html', event=event)

@mod.route('/<slug>/edit/', methods=['GET', 'POST'])
@login_required
def edit_event(slug):
    event = get_object_or_404(Event, Event.slug == slug)
    form = EditEventForm(
            CombinedMultiDict((request.form, request.files)),
            event
            )

    if form.validate_on_submit():
        event.title = form.title.data
        event.content = form.content.data
        event.published = form.published.data
        event.start_time = form.start_time.data
        event.location = form.location.data
        if form.upload.has_file():
            event.image = images.save(form.upload.data)
        event.save()
        return redirect(url_for('.view_event', slug=event.slug))

    return render_template('events/edit-event.html', form=form)

@mod.route('/<slug>/remove/')
@login_required
def remove_event(slug):
    event = get_object_or_404(Event, Event.slug == slug)
    event.delete_instance()
    return redirect(url_for('.overview'))
<|endoftext|>


<commit_before>import tensorflow as tf
import numpy as np
from param import Param
from .model import GPModel
import transforms
import conditionals
from .mean_functions import Zero
from tf_hacks import eye

class SVGP(GPModel):
    """
    This is the Sparse Variational GP (SVGP). The key reference is

    @inproceedings{hensman2014scalable,
      title={Scalable Variational Gaussian Process Classification},
      author={Hensman, James and Matthews, Alexander G. de G. and Ghahramani, Zoubin},
      booktitle={Proceedings of AISTATS},
      year={2015}
    }

    """
    def __init__(self, X, Y, kern, likelihood, Z, mean_function=Zero(), num_latent=None, q_diag=False, whiten=True):
        """
        X is a data matrix, size N x D
        Y is a data matrix, size N x R
        kern, likelihood, mean_function are appropriate GPflow objects
        Z is a matrix of pseudo inputs, size M x D
        num_latent is the number of latent process to use, default to Y.shape[1]
        q_diag is a boolean. If True, the covariance is approximated by a diagonal matrix.
        whiten is a boolean. It True, we use the whitened represenation of the inducing points.
        """
        GPModel.__init__(self, X, Y, kern, likelihood, mean_function)
        self.q_diag, self.whiten = q_diag, whiten
        self.Z = Param(Z)
        self.num_latent = num_latent or Y.shape[1]
        self.num_inducing = Z.shape[0]

        self.q_mu = Param(np.zeros((self.num_inducing, self.num_latent)))
        if self.q_diag:
            self.q_sqrt = Param(np.ones((self.num_inducing, self.num_latent)), transforms.positive)
        else:
            self.q_sqrt = Param(np.array([np.eye(self.num_inducing) for _ in range(self.num_latent)]).swapaxes(0,2))

    def build_likelihood(self):
        """
        This gives a variational bound on the model likelihood.

        There are four possible cases here, all combinations of True/False for self.whiten and self.q_diag.
        """
        if self.whiten:
            #First compute KL[q(v) || p(v)]] for each column of v
            KL = 0.5*tf.reduce_sum(tf.square(self.q_mu)) - 0.5*self.num_inducing*self.num_latent
            if self.q_diag:
                KL += -tf.reduce_sum(tf.log(self.q_sqrt)) + 0.5*tf.reduce_sum(tf.square(self.q_sqrt))
            else:
                #here we loop through all the independent columns, extracting the triangular part.
                def f(w):
                    R = tf.triu(w)
                    return -tf.reduce_sum(tf.log(tf.user_ops.get_diag(R))) + 0.5*tf.reduce_sum(tf.square(R))
                KL += theano.scan(f, self.q_sqrt.swapaxes(0,2))[0].sum()
            fmean, fvar = conditionals.gaussian_gp_predict_whitened(self.X, self.Z, self.kern, self.q_mu, self.q_sqrt)
        else:
            L = tf.cholesky(self.kern.K(self.Z) + eye(self.num_inducing) * 1e-4)
            alpha = tf.user_ops.triangular_solve(L, self.q_mu, 'lower')
            KL = 0.5*tf.reduce_sum(tf.square(alpha)) - 0.5*self.num_inducing*self.num_latent +\
                 self.num_latent * tf.reduce_sum(tf.log(tf.user_ops.get_diag(L)))
            L_inv = tf.user_ops.triangular_solve(L, eye(self.num_inducing), 'lower')
            K_inv = tf.user_ops.triangular_solve(tf.transpose(L), L_inv, 'upper')
            if self.q_diag:
                KL += -tf.reduce_sum(tf.log(self.q_sqrt))
                KL += 0.5 * tf.reduce_sum(tf.expand_dims(tf.user_ops.get_diag(K_inv), 1) * tf.square(self.q_sqrt))
            else:
                def f(w):
                    R = tf.triu(w)
                    S_q = tf.matmul(tf.transpose(R), R)
                    return -tf.reduce_sum(tf.log(tf.user_ops.get_diag(R))) + 0.5*tf.reduce_sum(S_q * K_inv)
                KL += theano.scan(f, self.q_sqrt.swapaxes(0,2))[0].sum()
            fmean, fvar = conditionals.gaussian_gp_predict(self.X, self.Z, self.kern, self.q_mu, self.q_sqrt)

        #add in mean function:
        fmean += self.mean_function(self.X)
        return tf.reduce_sum(self.likelihood.variational_expectations(fmean, fvar, self.Y)) - KL

    def build_predict(self, Xnew):
        if self.whiten:
            mu, var =  conditionals.gaussian_gp_predict_whitened(Xnew, self.Z, self.kern, self.q_mu, self.q_sqrt)
        else:
            mu, var =  conditionals.gaussian_gp_predict(Xnew, self.Z, self.kern, self.q_mu, self.q_sqrt)
        return mu + self.mean_function(Xnew), var

<commit_msg>loop explication in svgp<commit_after>import tensorflow as tf
import numpy as np
from param import Param
from .model import GPModel
import transforms
import conditionals
from .mean_functions import Zero
from tf_hacks import eye

class SVGP(GPModel):
    """
    This is the Sparse Variational GP (SVGP). The key reference is

    @inproceedings{hensman2014scalable,
      title={Scalable Variational Gaussian Process Classification},
      author={Hensman, James and Matthews, Alexander G. de G. and Ghahramani, Zoubin},
      booktitle={Proceedings of AISTATS},
      year={2015}
    }

    """
    def __init__(self, X, Y, kern, likelihood, Z, mean_function=Zero(), num_latent=None, q_diag=False, whiten=True):
        """
        X is a data matrix, size N x D
        Y is a data matrix, size N x R
        kern, likelihood, mean_function are appropriate GPflow objects
        Z is a matrix of pseudo inputs, size M x D
        num_latent is the number of latent process to use, default to Y.shape[1]
        q_diag is a boolean. If True, the covariance is approximated by a diagonal matrix.
        whiten is a boolean. It True, we use the whitened represenation of the inducing points.
        """
        GPModel.__init__(self, X, Y, kern, likelihood, mean_function)
        self.q_diag, self.whiten = q_diag, whiten
        self.Z = Param(Z)
        self.num_latent = num_latent or Y.shape[1]
        self.num_inducing = Z.shape[0]

        self.q_mu = Param(np.zeros((self.num_inducing, self.num_latent)))
        if self.q_diag:
            self.q_sqrt = Param(np.ones((self.num_inducing, self.num_latent)), transforms.positive)
        else:
            self.q_sqrt = Param(np.array([np.eye(self.num_inducing) for _ in range(self.num_latent)]).swapaxes(0,2))

    def build_likelihood(self):
        """
        This gives a variational bound on the model likelihood.

        There are four possible cases here, all combinations of True/False for self.whiten and self.q_diag.
        """
        if self.whiten:
            #First compute KL[q(v) || p(v)]] for each column of v
            KL = 0.5*tf.reduce_sum(tf.square(self.q_mu)) - 0.5*self.num_inducing*self.num_latent
            if self.q_diag:
                KL += -tf.reduce_sum(tf.log(self.q_sqrt)) + 0.5*tf.reduce_sum(tf.square(self.q_sqrt))
            else:
                #here we loop through all the independent columns, extracting the triangular part.
                for d in range(self.num_latent):
                    L = tf.user_ops.triangle(self.q_sqrt[d,:,:], 'lower')
                    KL += -tf.reduce_sum(tf.log(tf.user_ops.get_diag(L))) + 0.5*tf.reduce_sum(tf.square(L))
            fmean, fvar = conditionals.gaussian_gp_predict_whitened(self.X, self.Z, self.kern, self.q_mu, self.q_sqrt)
        else:
            L = tf.cholesky(self.kern.K(self.Z) + eye(self.num_inducing) * 1e-4)
            alpha = tf.user_ops.triangular_solve(L, self.q_mu, 'lower')
            KL = 0.5*tf.reduce_sum(tf.square(alpha)) - 0.5*self.num_inducing*self.num_latent +\
                 self.num_latent * tf.reduce_sum(tf.log(tf.user_ops.get_diag(L)))
            L_inv = tf.user_ops.triangular_solve(L, eye(self.num_inducing), 'lower')
            K_inv = tf.user_ops.triangular_solve(tf.transpose(L), L_inv, 'upper')
            if self.q_diag:
                KL += -tf.reduce_sum(tf.log(self.q_sqrt))
                KL += 0.5 * tf.reduce_sum(tf.expand_dims(tf.user_ops.get_diag(K_inv), 1) * tf.square(self.q_sqrt))
            else:
                def f(w):
                    R = tf.triu(w)
                    S_q = tf.matmul(tf.transpose(R), R)
                    return -tf.reduce_sum(tf.log(tf.user_ops.get_diag(R))) + 0.5*tf.reduce_sum(S_q * K_inv)
                KL += theano.scan(f, self.q_sqrt.swapaxes(0,2))[0].sum()
            fmean, fvar = conditionals.gaussian_gp_predict(self.X, self.Z, self.kern, self.q_mu, self.q_sqrt)

        #add in mean function:
        fmean += self.mean_function(self.X)
        return tf.reduce_sum(self.likelihood.variational_expectations(fmean, fvar, self.Y)) - KL

    def build_predict(self, Xnew):
        if self.whiten:
            mu, var =  conditionals.gaussian_gp_predict_whitened(Xnew, self.Z, self.kern, self.q_mu, self.q_sqrt)
        else:
            mu, var =  conditionals.gaussian_gp_predict(Xnew, self.Z, self.kern, self.q_mu, self.q_sqrt)
        return mu + self.mean_function(Xnew), var

<|endoftext|>


<commit_before>"""A webapp2 server for ChunaaviTukka."""

import os
import urllib
import json
import logging

from google.appengine.api import users
from google.appengine.ext import ndb

from models import CTUser
from models import CTCandidate
from models import CTConstituency

import jinja2
import webapp2

JINJA_ENVIRONMENT = jinja2.Environment(
    loader=jinja2.FileSystemLoader(os.path.dirname(__file__)),
    extensions=['jinja2.ext.autoescape'],
    autoescape=True)

def user_setup(self):
    '''Function to do basic user URL setup'''
    google_user = users.get_current_user()
    if google_user:
        url = users.create_logout_url(self.request.uri)
        url_linktext = 'Logout'
        ct_user = CTUser.get_or_insert(google_user)
    else:
        url = users.create_login_url(self.request.uri)
        url_linktext = 'Login'
        ct_user = None
    return (ct_user, url, url_linktext)

def get_predictions(user_id):
    return [{'cons':{'name':'Varanasi','slug':'varanasi-up'},'candidate':{'name':'M.M.Malaviya','party':'BHU','coalition': 'ABC'}},{'cons':{'name':'Amritsar','slug':'amritsar-pu'},'candidate':{'name':'Guru Gobind Singh','party':'Gurdwara','coalition': 'XYZ'}}]

def get_constituency_info(contest_slug):
    if contest_slug != 'varanasi':
        return {'name':'Some city', 'state': 'some state', 'predictions':[{'candidate':{'name':'M.M.Malaviya','party':'BHU','coalition': 'ABC'},'support':100},{'candidate':{'name':'Guru Gobind Singh','party':'Gurdwara','coalition': 'XYZ'},'support':75}]}
    else:
        conskey = ndb.Key(CTConstituency, contest_slug)
        logging.error(conskey)
        #TODO exception handling
        cons = conskey.get()
        logging.error(cons)
        predictions = []
        for candidate_key in cons.candidates:
            c = candidate_key.get()
            predictions.append({'candidate':{'name':c.name,'party':c.party,'coalition': c.coalition},'support':get_support(conskey,candidate_key)})
        return {'name':cons.name, 'state': cons.state, 'predictions':predictions};

def get_support(conskey,candidate_key):
    '''overall support for a candidate in a constituency'''
    #TODO implement this
    return 100

class HomeHandler(webapp2.RequestHandler):
    '''Shows the home page'''
    def get(self):
        (ct_user, url, url_linktext) = user_setup(self)
        template_values = {
            'greetings': [],
            'url': url,
            'url_linktext': url_linktext,
        }

        self.response.headers['Content-Type'] = 'text/html'
        template = JINJA_ENVIRONMENT.get_template('templates/index.html')
        self.response.write(template.render(template_values))

class ContestPageHandler(webapp2.RequestHandler):
    '''Handler for showing a contest's page'''
    def get(self, contest_slug):
        cons_info = get_constituency_info(contest_slug)
        (ct_user, url, url_linktext) = user_setup(self)
        format = self.request.get("f")
        template_values = {
            'contest_info': cons_info,
            'url': url,
            'url_linktext': url_linktext,
        }

        if (format == 'json'):
          self.response.headers['Content-Type'] = 'application/json'
          json.dump(cons_info,self.response.out)
        else:
          self.response.headers['Content-Type'] = 'text/html'
          template = JINJA_ENVIRONMENT.get_template('templates/contest.html')
          self.response.write(template.render(template_values))

class UserPageHandler(webapp2.RequestHandler):
    '''Handler for showing a user's page'''
    def get(self, user_id):
        predictions = get_predictions(user_id)
        (ct_user, url, url_linktext) = user_setup(self)

        format = self.request.get("f")
        template_values = {
            'user_name': user_id,
            'url': url,
            'url_linktext': url_linktext,
            'predictions': predictions,
        }

        if (format == 'json'):
          self.response.headers['Content-Type'] = 'application/json'
          json.dump(predictions,self.response.out)
        else:
          self.response.headers['Content-Type'] = 'text/html'
          template = JINJA_ENVIRONMENT.get_template('templates/user.html')
          self.response.write(template.render(template_values))

class UserPredictionHandler(webapp2.RequestHandler):
    '''Handler for recording user predictions'''
    def get(self, user_id, contest_slug):
        self.response.headers['Content-Type'] = 'text/plain'
        self.response.write('Show contest detail!')

class SettingsPageHandler(webapp2.RequestHandler):
    '''Show and manage settings page'''
    def get(self):
        (ct_user, url, url_linktext) = user_setup(self)
        if not ct_user:
            return webapp2.redirect(url)
        template_values = {
            'ct_user': ct_user,
            'url': url,
            'url_linktext': url_linktext,
            'error_message':'',
        }
        self.response.headers['Content-Type'] = 'text/html'
        template = JINJA_ENVIRONMENT.get_template('templates/settings.html')
        self.response.write(template.render(template_values))
    def post(self):
        (ct_user, url, url_linktext) = user_setup(self)
        if not ct_user:
            return webapp2.redirect(url)
        user_dn = self.request.get("display_name")
        error_message = ''

        if (ct_user.display_name == user_dn):
            update_success = False
            error_message = 'New name same as old.'
        elif self.is_display_name_disallowed(user_dn):
            update_success = False
            error_message = 'New name not allowed.'
        elif CTUser.is_display_name_taken(user_dn):
            update_success = False
            error_message = 'Someone already has this name.'
        else:
            ct_user.display_name = user_dn
            ct_user.put()
            update_success = True

        template_values = {
            'success': update_success,
            'error_message': error_message,
            'ct_user': ct_user,
            'url': url,
            'url_linktext': url_linktext,
        }
        self.response.headers['Content-Type'] = 'text/html'
        template = JINJA_ENVIRONMENT.get_template('templates/settings.html')
        self.response.write(template.render(template_values))

    def is_display_name_disallowed(self,user_dn):
        invalid_users = set(['about','account','admin','administrator','administration','app','api','backup','bin','bot','bots','cache','chi','config','db','dev','download','edit','forum','feed','faq','ftp','help','home','index','login','logout','php','public','settings','system','task','username','xxx','you'])
        return user_dn.lower() in invalid_users

class TempAddHandler(webapp2.RequestHandler):
    '''Show and manage settings page'''
    def get(self):
        candidate1 = CTCandidate(name='Arvind Kejriwal',party='AAP')
        candidate1.put()
        candidate2 = CTCandidate(name='Narendra Modi',party='BJP', coalition='NDA')
        candidate2.put()
        candidate3 = CTCandidate(name='Ajay Rai',party='Congress', coalition='UPA')
        candidate3.put()
        candidate4 = CTCandidate(name='Vijay Jaiswal',party='BSP')
        candidate4.put()
        candidate5 = CTCandidate(name='Mukhtar Ansari',party='Independent')
        candidate5.put()
        candidate6 = CTCandidate(name='Kailash Chaurasia',party='SP')
        candidate6.put()

        conskey = ndb.Key(CTConstituency, 'varanasi')
        cons = CTConstituency(key=conskey,name='Varanasi',state='Uttar Pradesh', candidates=[candidate1.key,candidate2.key,candidate3.key,candidate4.key,candidate5.key])
        cons.put()

application = webapp2.WSGIApplication([
    webapp2.Route(r'/', handler=HomeHandler, name='home'),
    webapp2.Route(r'/c/<contest_slug>/', handler=ContestPageHandler, name='contest_page'),
    webapp2.Route(r'/s/', handler=SettingsPageHandler, name='settings_page'),
    webapp2.Route(r'/u/<user_id>/', handler=UserPageHandler, name='user_page'),
    webapp2.Route(r'/u/<user_id>/<contest_slug>/', handler=UserPredictionHandler, name='user_prediction'),
#    webapp2.Route(r'/adddata/', handler=TempAddHandler, name='temp_addition'),
], debug=True)

# URL mapping, for reference
# /c/contest_slug - a contest's public page
# returns: list of public predictions about the contest, including summaries
# /u/user-id - a user's public page
# returns: list of public predictions by that user
# /u/user-id/contest-slug - user u's prediction about contest c
# returns: JSON object with candidate id, name, and party
<commit_msg>Resolves #4<commit_after>"""A webapp2 server for ChunaaviTukka."""

import os
import urllib
import json
import logging
import re

from google.appengine.api import users
from google.appengine.ext import ndb

from models import CTUser
from models import CTCandidate
from models import CTConstituency

import jinja2
import webapp2

JINJA_ENVIRONMENT = jinja2.Environment(
    loader=jinja2.FileSystemLoader(os.path.dirname(__file__)),
    extensions=['jinja2.ext.autoescape'],
    autoescape=True)

def user_setup(self):
    '''Function to do basic user URL setup'''
    google_user = users.get_current_user()
    if google_user:
        url = users.create_logout_url(self.request.uri)
        url_linktext = 'Logout'
        ct_user = CTUser.get_or_insert(google_user)
    else:
        url = users.create_login_url(self.request.uri)
        url_linktext = 'Login'
        ct_user = None
    return (ct_user, url, url_linktext)

def get_predictions(user_id):
    return [{'cons':{'name':'Varanasi','slug':'varanasi-up'},'candidate':{'name':'M.M.Malaviya','party':'BHU','coalition': 'ABC'}},{'cons':{'name':'Amritsar','slug':'amritsar-pu'},'candidate':{'name':'Guru Gobind Singh','party':'Gurdwara','coalition': 'XYZ'}}]

def get_constituency_info(contest_slug):
    if contest_slug != 'varanasi':
        return {'name':'Some city', 'state': 'some state', 'predictions':[{'candidate':{'name':'M.M.Malaviya','party':'BHU','coalition': 'ABC'},'support':100},{'candidate':{'name':'Guru Gobind Singh','party':'Gurdwara','coalition': 'XYZ'},'support':75}]}
    else:
        conskey = ndb.Key(CTConstituency, contest_slug)
        logging.error(conskey)
        #TODO exception handling
        cons = conskey.get()
        logging.error(cons)
        predictions = []
        for candidate_key in cons.candidates:
            c = candidate_key.get()
            predictions.append({'candidate':{'name':c.name,'party':c.party,'coalition': c.coalition},'support':get_support(conskey,candidate_key)})
        return {'name':cons.name, 'state': cons.state, 'predictions':predictions};

def get_support(conskey,candidate_key):
    '''overall support for a candidate in a constituency'''
    #TODO implement this
    return 100

class HomeHandler(webapp2.RequestHandler):
    '''Shows the home page'''
    def get(self):
        (ct_user, url, url_linktext) = user_setup(self)
        template_values = {
            'greetings': [],
            'url': url,
            'url_linktext': url_linktext,
        }

        self.response.headers['Content-Type'] = 'text/html'
        template = JINJA_ENVIRONMENT.get_template('templates/index.html')
        self.response.write(template.render(template_values))

class ContestPageHandler(webapp2.RequestHandler):
    '''Handler for showing a contest's page'''
    def get(self, contest_slug):
        cons_info = get_constituency_info(contest_slug)
        (ct_user, url, url_linktext) = user_setup(self)
        format = self.request.get("f")
        template_values = {
            'contest_info': cons_info,
            'url': url,
            'url_linktext': url_linktext,
        }

        if (format == 'json'):
          self.response.headers['Content-Type'] = 'application/json'
          json.dump(cons_info,self.response.out)
        else:
          self.response.headers['Content-Type'] = 'text/html'
          template = JINJA_ENVIRONMENT.get_template('templates/contest.html')
          self.response.write(template.render(template_values))

class UserPageHandler(webapp2.RequestHandler):
    '''Handler for showing a user's page'''
    def get(self, user_id):
        predictions = get_predictions(user_id)
        (ct_user, url, url_linktext) = user_setup(self)

        format = self.request.get("f")
        template_values = {
            'user_name': user_id,
            'url': url,
            'url_linktext': url_linktext,
            'predictions': predictions,
        }

        if (format == 'json'):
          self.response.headers['Content-Type'] = 'application/json'
          json.dump(predictions,self.response.out)
        else:
          self.response.headers['Content-Type'] = 'text/html'
          template = JINJA_ENVIRONMENT.get_template('templates/user.html')
          self.response.write(template.render(template_values))

class UserPredictionHandler(webapp2.RequestHandler):
    '''Handler for recording user predictions'''
    def get(self, user_id, contest_slug):
        self.response.headers['Content-Type'] = 'text/plain'
        self.response.write('Show contest detail!')

class SettingsPageHandler(webapp2.RequestHandler):
    '''Show and manage settings page'''
    def get(self):
        (ct_user, url, url_linktext) = user_setup(self)
        if not ct_user:
            return webapp2.redirect(url)
        template_values = {
            'ct_user': ct_user,
            'url': url,
            'url_linktext': url_linktext,
            'error_message':'',
        }
        self.response.headers['Content-Type'] = 'text/html'
        template = JINJA_ENVIRONMENT.get_template('templates/settings.html')
        self.response.write(template.render(template_values))
    def post(self):
        (ct_user, url, url_linktext) = user_setup(self)
        if not ct_user:
            return webapp2.redirect(url)
        user_dn = self.request.get("display_name")
        error_message = ''

        if (ct_user.display_name == user_dn):
            update_success = False
            error_message = 'New name same as old.'
        elif self.is_display_name_disallowed(user_dn):
            update_success = False
            error_message = 'New name not allowed.'
        elif CTUser.is_display_name_taken(user_dn):
            update_success = False
            error_message = 'Someone already has this name.'
        else:
            ct_user.display_name = user_dn
            ct_user.put()
            update_success = True

        template_values = {
            'success': update_success,
            'error_message': error_message,
            'ct_user': ct_user,
            'url': url,
            'url_linktext': url_linktext,
        }
        self.response.headers['Content-Type'] = 'text/html'
        template = JINJA_ENVIRONMENT.get_template('templates/settings.html')
        self.response.write(template.render(template_values))

    def is_display_name_disallowed(self,user_dn):
        invalid_users = set(['about','account','admin','administrator','administration','app','api','backup','bin','bot','bots','cache','chi','config','db','dev','download','edit','forum','feed','faq','ftp','help','home','index','login','logout','php','public','settings','system','task','username','xxx','you'])
        if user_dn.lower() in invalid_users :
            return True
        else:
            p = re.compile('^[a-zA-Z0-9#@_][a-zA-Z0-9#@_ ]*$') #chars,nums,and some others allowed, no spaces in beginning
            return not p.match(user_dn)

class TempAddHandler(webapp2.RequestHandler):
    '''Show and manage settings page'''
    def get(self):
        candidate1 = CTCandidate(name='Arvind Kejriwal',party='AAP')
        candidate1.put()
        candidate2 = CTCandidate(name='Narendra Modi',party='BJP', coalition='NDA')
        candidate2.put()
        candidate3 = CTCandidate(name='Ajay Rai',party='Congress', coalition='UPA')
        candidate3.put()
        candidate4 = CTCandidate(name='Vijay Jaiswal',party='BSP')
        candidate4.put()
        candidate5 = CTCandidate(name='Mukhtar Ansari',party='Independent')
        candidate5.put()
        candidate6 = CTCandidate(name='Kailash Chaurasia',party='SP')
        candidate6.put()

        conskey = ndb.Key(CTConstituency, 'varanasi')
        cons = CTConstituency(key=conskey,name='Varanasi',state='Uttar Pradesh', candidates=[candidate1.key,candidate2.key,candidate3.key,candidate4.key,candidate5.key])
        cons.put()

application = webapp2.WSGIApplication([
    webapp2.Route(r'/', handler=HomeHandler, name='home'),
    webapp2.Route(r'/c/<contest_slug>/', handler=ContestPageHandler, name='contest_page'),
    webapp2.Route(r'/s/', handler=SettingsPageHandler, name='settings_page'),
    webapp2.Route(r'/u/<user_id>/', handler=UserPageHandler, name='user_page'),
    webapp2.Route(r'/u/<user_id>/<contest_slug>/', handler=UserPredictionHandler, name='user_prediction'),
#    webapp2.Route(r'/adddata/', handler=TempAddHandler, name='temp_addition'),
], debug=True)

# URL mapping, for reference
# /c/contest_slug - a contest's public page
# returns: list of public predictions about the contest, including summaries
# /u/user-id - a user's public page
# returns: list of public predictions by that user
# /u/user-id/contest-slug - user u's prediction about contest c
# returns: JSON object with candidate id, name, and party
<|endoftext|>


<commit_before>from celery import shared_task
from celery import uuid as celery_uuid
import csv
import datetime
from django.conf import settings
import django.core.files
from django.core.files.storage import FileSystemStorage
from django.utils.text import slugify
import logging
import os.path
import requests
import tempfile
import time

from disease.files_utils import get_genome_filepath
from disease.models import AnalyzeDataOrder
from disease.tasks import recompute_genome_file
from twentythree.models import CeleryTask23
from twentythree.models import Token23

log = logging.getLogger('genoome.twentythree.tasks')

def twentythree_snp_mapping():
    snp_path = settings.SNP_MAPPING_FILEPATH23

    if not os.path.exists(snp_path):
        log.warning('Fetching snps mapping from 23andme.com')
        response = requests.get('https://api.23andme.com/res/txt/snps.b4e00fe1db50.data')
        with open(snp_path, 'w') as f:
            f.write(response.text)
        log.warning('New snps mapping23 saved at: %s', snp_path)

    with open(snp_path) as snp_mapping_f:
        csv_reader = csv.reader(snp_mapping_f, delimiter='\t')
        for row in csv_reader:
            if len(row) == 1:
                log.debug('Ignoring: %r', row)
                continue
            index, rsid, chromosome, position = row
            if not index.isdigit():
                log.debug('idx notdigit. Inoring: %r', row)
                continue
            index = int(index)
            yield index*2, rsid, chromosome, position

@shared_task
def fetch_genome_and_push_forward(ctask_pk):
    ctask = CeleryTask23.objects.get(pk=ctask_pk)
    ctask.status = 'fetching'
    ctask.save()
    def _may_raise():
        user_pk = ctask.user_id
        profile_id = ctask.chosen_profile

        filename = 'twentythree_%s_%s_%s.csv' % (user_pk, slugify(profile_id), int(time.time()))
        filepath = get_genome_filepath(ctask.user, filename)

        from django.core.files.storage import FileSystemStorage
        storage = FileSystemStorage()
        token = Token23.objects.get(user_id=user_pk)
        log.info('Fetching genome...')
        try:
            genome = token.get_genome(profile_id)
            log.info('Got genome len(%d)', len(genome))
        except Token23.ClientError as e:
            if e.json()['error'] == 'access_denied':
                log.warning('Got access_denied. Refreshing token:%s', token.pk)
                token.refresh()
            else:
                raise

        with tempfile.TemporaryFile('w+') as csv_file:
            log.info('Formatting csv for tempfile')
            csv_file.write('# 23andme\n') # for disease.files_utils.get_parser() to detect 23andme format
            csv_file.write('# generated by twentythree at UTC %s\n' % (datetime.datetime.utcnow(),))
            csv_file.write('# fetched from api\n')
            csv_writer = csv.writer(csv_file, delimiter='\t')
            for index, rsid, chromosome, position in twentythree_snp_mapping():
                log.debug('index: %r rsid: %r chromosome: %r position: %s', index, rsid, chromosome, position)
                if not rsid.startswith('rs'):
                    log.debug('ignoring: %r', rsid)
                    continue
                genotype = genome[index:index+2]
                log.debug('genotype: %r', genotype)
                csv_writer.writerow([rsid, chromosome, position, genotype])
            log.info('Saving csv to storage...')
            storage.save(filepath, django.core.files.File(csv_file))

        task_id = celery_uuid()
        analyze_order = AnalyzeDataOrder.objects.create(uploaded_filename=filename, user_id=user_pk, task_uuid=task_id)
        ctask.analyze_order = analyze_order
        ctask.status = 'parsing'
        ctask.save()
        log.info('Calling recompute_genome_file...')
        recompute_genome_file.apply_async(args=(filepath,), task_id=task_id)
        log.debug('Done.')
    try:
        _may_raise()
    except Exception as e:
        log.exception('Problem running fetch_genome_and_push_forward().')
        ctask.status = 'error'
        ctask.save()
        raise
<commit_msg>23andme wip<commit_after>from celery import shared_task
from celery import uuid as celery_uuid
import csv
import datetime
from django.conf import settings
import django.core.files
from django.core.files.storage import FileSystemStorage
from django.utils.text import slugify
import logging
import os.path
import requests
import tempfile
import time

from disease.files_utils import get_genome_filepath
from disease.models import AnalyzeDataOrder
from disease.tasks import recompute_genome_file
from twentythree.models import CeleryTask23
from twentythree.models import Token23

log = logging.getLogger('genoome.twentythree.tasks')

def twentythree_snp_mapping():
    snp_path = settings.SNP_MAPPING_FILEPATH23

    if not os.path.exists(snp_path):
        log.warning('Fetching snps mapping from 23andme.com')
        response = requests.get('https://api.23andme.com/res/txt/snps.b4e00fe1db50.data')
        with open(snp_path, 'w') as f:
            f.write(response.text)
        log.warning('New snps mapping23 saved at: %s', snp_path)

    with open(snp_path) as snp_mapping_f:
        csv_reader = csv.reader(snp_mapping_f, delimiter='\t')
        for row in csv_reader:
            if len(row) == 1:
                log.debug('Ignoring: %r', row)
                continue
            index, rsid, chromosome, position = row
            if not index.isdigit():
                log.debug('idx notdigit. Inoring: %r', row)
                continue
            index = int(index)
            yield index*2, rsid, chromosome, position

@shared_task
def fetch_genome_and_push_forward(ctask_pk):
    ctask = CeleryTask23.objects.get(pk=ctask_pk)
    ctask.status = 'fetching'
    ctask.save()
    def _may_raise():
        user_pk = ctask.user_id
        profile_id = ctask.chosen_profile

        filename = 'twentythree_%s_%s_%s.csv' % (user_pk, slugify(profile_id), int(time.time()))
        filepath = get_genome_filepath(ctask.user, filename)

        from django.core.files.storage import FileSystemStorage
        storage = FileSystemStorage()
        token = Token23.objects.get(user_id=user_pk)
        log.info('Fetching genome...')
        try:
            genome = token.get_genome(profile_id)
            log.info('Got genome len(%d)', len(genome))
        except Token23.ClientError as e:
            if e.json()['error'] == 'access_denied':
                log.warning('Got access_denied. Refreshing token:%s', token.pk)
                token.refresh()
            else:
                raise

        with tempfile.TemporaryFile('w+') as csv_file:
            log.info('Formatting csv for tempfile')
            csv_file.write('# 23andme\n') # for disease.files_utils.get_parser() to detect 23andme format
            csv_file.write('# generated by twentythree at UTC %s\n' % (datetime.datetime.utcnow(),))
            csv_file.write('# fetched from api\n')
            csv_writer = csv.writer(csv_file, delimiter='\t')
            for index, rsid, chromosome, position in twentythree_snp_mapping():
                #log.debug('index: %r rsid: %r chromosome: %r position: %s', index, rsid, chromosome, position)
                if not rsid.startswith('rs'):
                    log.debug('ignoring: %r', rsid)
                    continue
                genotype = genome[index:index+2]
                #log.debug('genotype: %r', genotype)
                csv_writer.writerow([rsid, chromosome, position, genotype])
            log.info('Saving csv to storage...')
            storage.save(filepath, django.core.files.File(csv_file))

        task_id = celery_uuid()
        analyze_order = AnalyzeDataOrder.objects.create(uploaded_filename=filename, user_id=user_pk, task_uuid=task_id)
        ctask.analyze_order = analyze_order
        ctask.status = 'parsing'
        ctask.save()
        log.info('Calling recompute_genome_file...')
        recompute_genome_file.apply_async(args=(filepath,), task_id=task_id)
        log.debug('Done.')
    try:
        _may_raise()
    except Exception as e:
        log.exception('Problem running fetch_genome_and_push_forward().')
        ctask.status = 'error'
        ctask.save()
        raise
<|endoftext|>


<commit_before>"""`appengine_config` gets loaded when starting a new application instance."""
from google.appengine.ext import vendor

# insert `lib` as a site directory so our `main` module can load
# third-party libraries, and override built-ins with newer
# versions.
vendor.add('lib')
<commit_msg>Updating documentation in appengine_config.py<commit_after>"""
`appengine_config.py` is automatically loaded when Google App Engine
starts a new instance of your application. This runs before any
WSGI applications specified in app.yaml are loaded.
"""

from google.appengine.ext import vendor

# Third-party libraries are stored in "lib", vendoring will make
# sure that they are importable by the application.
vendor.add('lib')
<|endoftext|>


<commit_before>"""All models for the extension."""

import sys
import json
import csv

from osgeo import ogr

from django.conf import settings
from django.dispatch import receiver
from django.db import models
from django.template.defaultfilters import slugify
from django.contrib.postgres.fields import ArrayField
from django.contrib.gis.db import models as gis

from django_pgjson.fields import JsonBField
from model_utils.models import StatusModel, TimeStampedModel

from geokey.projects.models import Project
from geokey.categories.models import Category, Field

from geokey_dataimports.helpers.model_helpers import import_from_csv
from .helpers import type_helpers
from .base import STATUS, FORMAT
from .exceptions import FileParseError
from .managers import DataImportManager

class DataImport(StatusModel, TimeStampedModel):
    """Store a single data import."""

    STATUS = STATUS
    FORMAT = FORMAT

    name = models.CharField(max_length=100)
    description = models.TextField(null=True, blank=True)
    dataformat = models.CharField(max_length=10, null=False, choices=FORMAT)
    file = models.FileField(
        upload_to='dataimports/files',
        max_length=500
    )
    keys = ArrayField(models.CharField(max_length=100), null=True, blank=True)

    project = models.ForeignKey(
        'projects.Project',
        related_name='dataimports'
    )
    category = models.ForeignKey(
        'categories.Category',
        null=True,
        blank=True
    )
    creator = models.ForeignKey(settings.AUTH_USER_MODEL)

    objects = DataImportManager()

    def delete(self, *args, **kwargs):
        """Delete the data import by setting its status to `deleted`."""
        self.status = self.STATUS.deleted
        self.save()

    def get_lookup_fields(self):
        """Get all lookup fields of a category."""
        lookupfields = {}
        for field in self.category.fields.all():
            if field.fieldtype == 'LookupField':
                lookupfields[field.key] = field
        return lookupfields

@receiver(models.signals.post_save, sender=DataImport)
def post_save_dataimport(sender, instance, created, **kwargs):
    """Map data fields and data features when the data import gets created."""
    if created:
        datafields = []
        datafeatures = []

        fields = []
        features = []
        errors = []

        if instance.dataformat == FORMAT.KML:
            driver = ogr.GetDriverByName('KML')
            reader = driver.Open(instance.file.path)

            for layer in reader:
                for feature in layer:
                    features.append(feature.ExportToJson())
        else:
            csv.field_size_limit(sys.maxsize)
            file_obj = open(instance.file.path, 'rU')

        if instance.dataformat == FORMAT.GeoJSON:
            reader = json.load(file)
            features = reader['features']

        if instance.dataformat == FORMAT.CSV:
            import_from_csv(features=features, fields=fields, file_obj=file_obj)

        for feature in features:
            geometries = {}

            for key, value in feature['properties'].iteritems():
                field = None

                for existing_field in fields:
                    if existing_field['name'] == key:
                        field = existing_field
                        break

                if field is None:
                    fields.append({
                        'name': key,
                        'good_types': set(['TextField', 'LookupField']),
                        'bad_types': set([])
                    })
                    field = fields[-1]

                fieldtype = None

                if 'geometry' not in feature:
                    try:
                        geometry = ogr.CreateGeometryFromWkt(str(value))
                        geometry = geometry.ExportToJson()
                    except:
                        geometry = None

                    fieldtype = 'GeometryField'
                    if geometry is not None:
                        if fieldtype not in field['bad_types']:
                            field['good_types'].add(fieldtype)
                            geometries[field['name']] = json.loads(geometry)
                    else:
                        field['good_types'].discard(fieldtype)
                        field['bad_types'].add(fieldtype)
                        fieldtype = None

                if fieldtype is None:
                    fieldtype = 'NumericField'
                    if type_helpers.is_numeric(value):
                        if fieldtype not in field['bad_types']:
                            field['good_types'].add(fieldtype)
                    else:
                        field['good_types'].discard(fieldtype)
                        field['bad_types'].add(fieldtype)

                    fieldtypes = ['DateField', 'DateTimeField']
                    if type_helpers.is_date(value):
                        for fieldtype in fieldtypes:
                            if fieldtype not in field['bad_types']:
                                field['good_types'].add(fieldtype)
                    else:
                        for fieldtype in fieldtypes:
                            field['good_types'].discard(fieldtype)
                            field['bad_types'].add(fieldtype)

                    fieldtype = 'TimeField'
                    if type_helpers.is_time(value):
                        if fieldtype not in field['bad_types']:
                            field['good_types'].add(fieldtype)
                    else:
                        field['good_types'].discard(fieldtype)
                        field['bad_types'].add(fieldtype)

            if 'geometry' not in feature and len(geometries) == 0:
                errors.append({
                    'line': feature['line'],
                    'messages': ['The entry has no geometry set.']
                })
            else:
                feature['geometries'] = geometries

        geometryfield = None
        for field in fields:
            if 'GeometryField' not in field['good_types']:
                datafields.append({
                    'name': field['name'],
                    'types': list(field['good_types'])
                })
            elif geometryfield is None:
                geometryfield = field['name']

        for feature in features:
            geometry = None
            if 'geometry' in feature:
                geometry = feature['geometry']
            elif 'geometries' in feature:
                if not geometryfield:
                    errors.append({
                        'line': feature['line'],
                        'messages': ['The file has no valid geometry field.']
                    })
                else:
                    geometries = feature['geometries']
                    if geometryfield in geometries:
                        geometry = geometries[geometryfield]

            if geometry:
                datafeatures.append({
                    'geometry': geometry,
                    'properties': feature['properties']
                })

        if errors:
            instance.delete()
            raise FileParseError('Failed to read file.', errors)
        else:
            for datafield in datafields:
                if datafield['name']:
                    DataField.objects.create(
                        name=datafield['name'],
                        types=list(datafield['types']),
                        dataimport=instance
                    )
            for datafeature in datafeatures:
                DataFeature.objects.create(
                    geometry=json.dumps(datafeature['geometry']),
                    properties=datafeature['properties'],
                    dataimport=instance
                )

class DataField(TimeStampedModel):
    """Store a single data field."""

    name = models.CharField(max_length=100)
    key = models.CharField(max_length=100, null=True, blank=True)
    types = ArrayField(models.CharField(max_length=100), null=True, blank=True)

    dataimport = models.ForeignKey(
        'DataImport',
        related_name='datafields'
    )

    def convert_to_field(self, name, fieldtype):
        """
        Convert data field to regular GeoKey field.

        Parameters
        ----------
        user : geokey.users.models.User
            The request user.
        name : str
            The name of the field.
        fieldtype : str
            The field type.

        Returns
        -------
        geokey.categories.models.Field
            The field created.
        """
        category = self.dataimport.category
        field = None

        if self.key:
            try:
                field = category.fields.get(key=self.key)
            except Category.DoesNotExist:
                pass

        proposed_key = slugify(self.name)
        suggested_key = proposed_key

        if field:
            suggested_key = field.key
        else:
            count = 1
            while category.fields.filter(key=suggested_key).exists():
                suggested_key = '%s-%s' % (proposed_key, count)
                count += 1

            self.key = suggested_key
            self.save()

            field = Field.create(
                name,
                self.key,
                '', False,
                category,
                fieldtype
            )

        if self.key != self.name:
            for datafeature in self.dataimport.datafeatures.all():
                properties = datafeature.properties
                if self.name in properties:
                    properties[self.key] = properties.pop(self.name)
                datafeature.properties = properties
                datafeature.save()

        return field

class DataFeature(TimeStampedModel):
    """Store a single data feature."""

    imported = models.BooleanField(default=False)
    geometry = gis.GeometryField(geography=True)
    properties = JsonBField(default={})

    dataimport = models.ForeignKey(
        'DataImport',
        related_name='datafeatures'
    )

@receiver(models.signals.post_save, sender=Project)
def post_save_project(sender, instance, **kwargs):
    """Remove associated data imports when the project gets deleted."""
    if instance.status == 'deleted':
        DataImport.objects.filter(project=instance).delete()

@receiver(models.signals.post_save, sender=Category)
def post_save_category(sender, instance, **kwargs):
    """Remove associated data imports when the category gets deleted."""
    if instance.status == 'deleted':
        DataImport.objects.filter(category=instance).delete()
<commit_msg>Fix loading JSON from file object<commit_after>"""All models for the extension."""

import sys
import json
import csv

from osgeo import ogr

from django.conf import settings
from django.dispatch import receiver
from django.db import models
from django.template.defaultfilters import slugify
from django.contrib.postgres.fields import ArrayField
from django.contrib.gis.db import models as gis

from django_pgjson.fields import JsonBField
from model_utils.models import StatusModel, TimeStampedModel

from geokey.projects.models import Project
from geokey.categories.models import Category, Field

from geokey_dataimports.helpers.model_helpers import import_from_csv
from .helpers import type_helpers
from .base import STATUS, FORMAT
from .exceptions import FileParseError
from .managers import DataImportManager

class DataImport(StatusModel, TimeStampedModel):
    """Store a single data import."""

    STATUS = STATUS
    FORMAT = FORMAT

    name = models.CharField(max_length=100)
    description = models.TextField(null=True, blank=True)
    dataformat = models.CharField(max_length=10, null=False, choices=FORMAT)
    file = models.FileField(
        upload_to='dataimports/files',
        max_length=500
    )
    keys = ArrayField(models.CharField(max_length=100), null=True, blank=True)

    project = models.ForeignKey(
        'projects.Project',
        related_name='dataimports'
    )
    category = models.ForeignKey(
        'categories.Category',
        null=True,
        blank=True
    )
    creator = models.ForeignKey(settings.AUTH_USER_MODEL)

    objects = DataImportManager()

    def delete(self, *args, **kwargs):
        """Delete the data import by setting its status to `deleted`."""
        self.status = self.STATUS.deleted
        self.save()

    def get_lookup_fields(self):
        """Get all lookup fields of a category."""
        lookupfields = {}
        for field in self.category.fields.all():
            if field.fieldtype == 'LookupField':
                lookupfields[field.key] = field
        return lookupfields

@receiver(models.signals.post_save, sender=DataImport)
def post_save_dataimport(sender, instance, created, **kwargs):
    """Map data fields and data features when the data import gets created."""
    if created:
        datafields = []
        datafeatures = []

        fields = []
        features = []
        errors = []

        if instance.dataformat == FORMAT.KML:
            driver = ogr.GetDriverByName('KML')
            reader = driver.Open(instance.file.path)

            for layer in reader:
                for feature in layer:
                    features.append(feature.ExportToJson())
        else:
            csv.field_size_limit(sys.maxsize)
            file_obj = open(instance.file.path, 'rU')

        if instance.dataformat == FORMAT.GeoJSON:
            reader = json.load(file_obj)
            features = reader['features']

        if instance.dataformat == FORMAT.CSV:
            import_from_csv(features=features, fields=fields, file_obj=file_obj)

        for feature in features:
            geometries = {}

            for key, value in feature['properties'].iteritems():
                field = None

                for existing_field in fields:
                    if existing_field['name'] == key:
                        field = existing_field
                        break

                if field is None:
                    fields.append({
                        'name': key,
                        'good_types': set(['TextField', 'LookupField']),
                        'bad_types': set([])
                    })
                    field = fields[-1]

                fieldtype = None

                if 'geometry' not in feature:
                    try:
                        geometry = ogr.CreateGeometryFromWkt(str(value))
                        geometry = geometry.ExportToJson()
                    except:
                        geometry = None

                    fieldtype = 'GeometryField'
                    if geometry is not None:
                        if fieldtype not in field['bad_types']:
                            field['good_types'].add(fieldtype)
                            geometries[field['name']] = json.loads(geometry)
                    else:
                        field['good_types'].discard(fieldtype)
                        field['bad_types'].add(fieldtype)
                        fieldtype = None

                if fieldtype is None:
                    fieldtype = 'NumericField'
                    if type_helpers.is_numeric(value):
                        if fieldtype not in field['bad_types']:
                            field['good_types'].add(fieldtype)
                    else:
                        field['good_types'].discard(fieldtype)
                        field['bad_types'].add(fieldtype)

                    fieldtypes = ['DateField', 'DateTimeField']
                    if type_helpers.is_date(value):
                        for fieldtype in fieldtypes:
                            if fieldtype not in field['bad_types']:
                                field['good_types'].add(fieldtype)
                    else:
                        for fieldtype in fieldtypes:
                            field['good_types'].discard(fieldtype)
                            field['bad_types'].add(fieldtype)

                    fieldtype = 'TimeField'
                    if type_helpers.is_time(value):
                        if fieldtype not in field['bad_types']:
                            field['good_types'].add(fieldtype)
                    else:
                        field['good_types'].discard(fieldtype)
                        field['bad_types'].add(fieldtype)

            if 'geometry' not in feature and len(geometries) == 0:
                errors.append({
                    'line': feature['line'],
                    'messages': ['The entry has no geometry set.']
                })
            else:
                feature['geometries'] = geometries

        geometryfield = None
        for field in fields:
            if 'GeometryField' not in field['good_types']:
                datafields.append({
                    'name': field['name'],
                    'types': list(field['good_types'])
                })
            elif geometryfield is None:
                geometryfield = field['name']

        for feature in features:
            geometry = None
            if 'geometry' in feature:
                geometry = feature['geometry']
            elif 'geometries' in feature:
                if not geometryfield:
                    errors.append({
                        'line': feature['line'],
                        'messages': ['The file has no valid geometry field.']
                    })
                else:
                    geometries = feature['geometries']
                    if geometryfield in geometries:
                        geometry = geometries[geometryfield]

            if geometry:
                datafeatures.append({
                    'geometry': geometry,
                    'properties': feature['properties']
                })

        if errors:
            instance.delete()
            raise FileParseError('Failed to read file.', errors)
        else:
            for datafield in datafields:
                if datafield['name']:
                    DataField.objects.create(
                        name=datafield['name'],
                        types=list(datafield['types']),
                        dataimport=instance
                    )
            for datafeature in datafeatures:
                DataFeature.objects.create(
                    geometry=json.dumps(datafeature['geometry']),
                    properties=datafeature['properties'],
                    dataimport=instance
                )

class DataField(TimeStampedModel):
    """Store a single data field."""

    name = models.CharField(max_length=100)
    key = models.CharField(max_length=100, null=True, blank=True)
    types = ArrayField(models.CharField(max_length=100), null=True, blank=True)

    dataimport = models.ForeignKey(
        'DataImport',
        related_name='datafields'
    )

    def convert_to_field(self, name, fieldtype):
        """
        Convert data field to regular GeoKey field.

        Parameters
        ----------
        user : geokey.users.models.User
            The request user.
        name : str
            The name of the field.
        fieldtype : str
            The field type.

        Returns
        -------
        geokey.categories.models.Field
            The field created.
        """
        category = self.dataimport.category
        field = None

        if self.key:
            try:
                field = category.fields.get(key=self.key)
            except Category.DoesNotExist:
                pass

        proposed_key = slugify(self.name)
        suggested_key = proposed_key

        if field:
            suggested_key = field.key
        else:
            count = 1
            while category.fields.filter(key=suggested_key).exists():
                suggested_key = '%s-%s' % (proposed_key, count)
                count += 1

            self.key = suggested_key
            self.save()

            field = Field.create(
                name,
                self.key,
                '', False,
                category,
                fieldtype
            )

        if self.key != self.name:
            for datafeature in self.dataimport.datafeatures.all():
                properties = datafeature.properties
                if self.name in properties:
                    properties[self.key] = properties.pop(self.name)
                datafeature.properties = properties
                datafeature.save()

        return field

class DataFeature(TimeStampedModel):
    """Store a single data feature."""

    imported = models.BooleanField(default=False)
    geometry = gis.GeometryField(geography=True)
    properties = JsonBField(default={})

    dataimport = models.ForeignKey(
        'DataImport',
        related_name='datafeatures'
    )

@receiver(models.signals.post_save, sender=Project)
def post_save_project(sender, instance, **kwargs):
    """Remove associated data imports when the project gets deleted."""
    if instance.status == 'deleted':
        DataImport.objects.filter(project=instance).delete()

@receiver(models.signals.post_save, sender=Category)
def post_save_category(sender, instance, **kwargs):
    """Remove associated data imports when the category gets deleted."""
    if instance.status == 'deleted':
        DataImport.objects.filter(category=instance).delete()
<|endoftext|>


<commit_before>import connexion

app = connexion.FlaskApp(__name__, specification_dir='api/')
app.add_api('getService.yaml')
app.run(port=8080)
<commit_msg>Updated app<commit_after>import connexion

app = connexion.App(__name__, 8080, specification_dir='api/')
app.add_api('getService.yaml', resolver=RestyResolver('api'))
app.run()
<|endoftext|>


<commit_before># This class implements the Websocket protocol draft version as of May 23, 2010
# The version as of August 6, 2010 will be implementend once Firefox or
# Webkit-trunk support this version.

import binascii
import struct

class WebSocket(object):
    def __init__(self, sock, rfile, environ):
        self.rfile = rfile
        self.socket = sock
        self.origin = environ.get('HTTP_ORIGIN')
        self.protocol = environ.get('HTTP_SEC_WEBSOCKET_PROTOCOL', 'unknown')
        self.path = environ.get('PATH_INFO')
        self.websocket_closed = False

    def send(self, message):
        if self.websocket_closed:
            raise Exception("Connection was terminated")

        if isinstance(message, unicode):
            message = message.encode('utf-8')
        elif isinstance(message, str):
            message = unicode(message).encode('utf-8')
        else:
            raise Exception("Invalid message encoding")

        self.socket.sendall("\x00" + message + "\xFF")

    def close_connection(self):
        if not self.websocket_closed:
            self.websocket_closed = True
            self.socket.shutdown(True)
            self.socket.close()
        else:
            return

    def _message_length(self):
        # TODO: buildin security agains lengths greater than 2**31 or 2**32
        length = 0

        while True:
            byte_str = self.rfile.read(1)

            if not byte_str:
                return 0
            else:
                byte = ord(byte_str)

            if byte != 0x00:
                length = length * 128 + (byte & 0x7f)
                if (byte & 0x80) != 0x80:
                    break

        return length

    def _read_until(self):
        bytes = []

        while True:
            byte = self.rfile.read(1)
            if ord(byte) != 0xff:
                bytes.append(byte)
            else:
                break

        return ''.join(bytes)

    def wait(self):
        while True:
            if self.websocket_closed:
                return None

            frame_str = self.rfile.read(1)
            if not frame_str:
                # Connection lost?
                self.websocket_closed = True
                continue
            else:
                frame_type = ord(frame_str)

            if (frame_type & 0x80) == 0x00: # most significant byte is not set

                if frame_type == 0x00:
                    bytes = self._read_until()
                    return bytes.decode("utf-8", "replace")
                else:
                    self.websocket_closed = True

            elif (frame_type & 0x80) == 0x80: # most significant byte is set
                # Read binary data (forward-compatibility)
                if frame_type != 0xff:
                    self.websocket_closed = True
                else:
                    length = self._message_length()
                    if length == 0:
                        self.websocket_closed = True
                    else:
                        self.rfile.read(length) # discard the bytes
            else:
                raise IOError("Reveiced an invalid message")

class WebSocketVersion7(WebSocket):
    FIN = int("10000000", 2)
    RSV = int("01110000", 2)
    OPCODE = int("00001111", 2)
    MASK = int("10000000", 2)
    PAYLOAD = int("01111111", 2)

    OPCODE_FRAG = 0x0
    OPCODE_TEXT = 0x1
    OPCODE_BINARY = 0x2
    OPCODE_CLOSE = 0x8
    OPCODE_PING = 0x9
    OPCODE_PONG = 0xA

    REASON_NORMAL = 1000
    REASON_GOING_AWAY = 1001
    REASON_PROTOCOL_ERROR = 1002

    LEN_16 = 126
    LEN_64 = 127

    def __init__(self, sock, rfile, environ):
        self.rfile = rfile
        self.socket = sock
        self.origin = environ.get('HTTP_SEC_WEBSOCKET_ORIGIN')
        self.protocol = environ.get('HTTP_SEC_WEBSOCKET_PROTOCOL', 'unknown')
        self.path = environ.get('PATH_INFO')
        self.websocket_closed = False
        self._fragments = []

    def _read_from_socket(self, count):
        return self.rfile.read(count)

    # TODO: replace all magic numbers with constants
    def wait(self):
        while True:
            payload = ''
            if self.websocket_closed:
                return None

            opcode_octet, length_octet = struct.unpack('!BB', self._read_from_socket(2))

            if self.RSV & opcode_octet:
                self.close(1002, 'Reserved bits cannot be set')
                return None

            opcode = opcode_octet & self.OPCODE
            is_final_frag = (self.FIN & opcode_octet) != 0

            if self._is_opcode_invalid(opcode):
                self.close(1002, 'Invalid opcode %x' % opcode)
                return None

            if not is_final_frag and self.OPCODE_CLOSE <= opcode <= self.OPCODE_PONG:
                self.close(self.REASON_PROTOCOL_ERROR, 'Control frames cannot be fragmented')
                return None

            if len(self._fragments) > 0 and not is_final_frag and opcode != self.OPCODE_FRAG:
                self.close(self.REASON_PROTOCOL_ERROR,
                        'Received new fragment frame with non-zero opcode')
                return None

            if len(self._fragments) > 0 and is_final_frag and (
                    self.OPCODE_TEXT <= opcode <= self.OPCODE_BINARY):
                self.close(self.REASON_PROTOCOL_ERROR,
                        'Received new unfragmented data frame during fragmented message')
                return None

            if not self.MASK & length_octet:
                self.close(1002, 'MASK must be set')
                return None

            length_code = length_octet & self.PAYLOAD

            if length_code > 125 and (self.OPCODE_CLOSE <= opcode <= self.OPCODE_PONG):
                self.close(1002, 'Control frame payload cannot be larger than 125 bytes')
                return None

            if length_code < 126:
                length = length_code
            elif length_code == 126:
                length = struct.unpack('!H', self._read_from_socket(2))[0]
            elif length_code == 127:
                length = struct.unpack('!Q', self._read_from_socket(8))[0]
            else:
                raise Exception('Calculated invalid length')

            mask_octets = struct.unpack('!BBBB', self._read_from_socket(4))
            masked_payload = self._read_from_socket(length)

            payload = ''

            j = 0
            for c in masked_payload:
                # TODO: optimize me? http://www.skymind.com/~ocrow/python_string/
                payload += chr(ord(c) ^ mask_octets[j])
                j = (j + 1) % 4

            if opcode == self.OPCODE_TEXT:
                payload = payload.decode('utf-8')
            elif opcode == self.OPCODE_CLOSE:
                if length >= 2:
                    reason, message = struct.unpack('!H%ds' % (length - 2), payload)
                else:
                    reason = message = None

                self.close(1000, '')
                return (reason, message)

            if opcode == self.OPCODE_PING:
                self.send(self.OPCODE_PONG, payload)
                return (self.OPCODE_PING, payload)
            elif opcode == self.OPCODE_PONG:
                return (self.OPCODE_PONG, payload)

            if is_final_frag:
                payload = ''.join(self._fragments) + payload
                self._fragments = []
                return payload
            else:
                self._fragments.append(payload)

    def _encode_text(self, s):
        if isinstance(s, unicode):
            return s.encode('utf-8')
        elif isinstance(s, str):
            return unicode(s).encode('utf-8')
        else:
            raise Exception('Invalid encoding')

    def _is_opcode_invalid(self, opcode):
        return opcode < self.OPCODE_FRAG or (opcode > self.OPCODE_BINARY and
                opcode < self.OPCODE_CLOSE) or opcode > self.OPCODE_PONG

    def send(self, opcode, message):
        if self.websocket_closed:
            raise Exception('Connection was terminated')

        if self._is_opcode_invalid(opcode):
            raise Exception('Invalid opcode %d' % opcode)

        if opcode == self.OPCODE_TEXT:
            message = self._encode_text(message)

        length = len(message)

        if opcode == self.OPCODE_TEXT:
            message = struct.pack('!%ds' % length, message)

        if length < 126:
            preamble = struct.pack('!BB', self.FIN | opcode, length)
        elif length < 2 ** 16:
            preamble = struct.pack('!BBH', self.FIN | opcode, self.LEN_16, length)
        elif length < 2 ** 64:
            preamble = struct.pack('!BBQ', self.FIN | opcode, self.LEN_64, length)
        else:
            # this can't really happen, but for correctness sake...
            raise Exception('Message is too long')

        self.socket.sendall(preamble + message)

    def close(self, reason, message):
        message = self._encode_text(message)
        self.send(self.OPCODE_CLOSE, struct.pack('!H%ds' % len(message), reason, message))
        self.websocket_closed = True

        # based on gevent/pywsgi.py
        # see http://pypi.python.org/pypi/gevent#downloads
        if self.socket is not None:
            try:
                self.socket._sock.close()
                self.socket.close()
            except socket.error:
                pass
<commit_msg>Add remaining reason constants and replace magic numbers<commit_after># This class implements the Websocket protocol draft version as of May 23, 2010
# The version as of August 6, 2010 will be implementend once Firefox or
# Webkit-trunk support this version.

import binascii
import struct

class WebSocket(object):
    def __init__(self, sock, rfile, environ):
        self.rfile = rfile
        self.socket = sock
        self.origin = environ.get('HTTP_ORIGIN')
        self.protocol = environ.get('HTTP_SEC_WEBSOCKET_PROTOCOL', 'unknown')
        self.path = environ.get('PATH_INFO')
        self.websocket_closed = False

    def send(self, message):
        if self.websocket_closed:
            raise Exception("Connection was terminated")

        if isinstance(message, unicode):
            message = message.encode('utf-8')
        elif isinstance(message, str):
            message = unicode(message).encode('utf-8')
        else:
            raise Exception("Invalid message encoding")

        self.socket.sendall("\x00" + message + "\xFF")

    def close_connection(self):
        if not self.websocket_closed:
            self.websocket_closed = True
            self.socket.shutdown(True)
            self.socket.close()
        else:
            return

    def _message_length(self):
        # TODO: buildin security agains lengths greater than 2**31 or 2**32
        length = 0

        while True:
            byte_str = self.rfile.read(1)

            if not byte_str:
                return 0
            else:
                byte = ord(byte_str)

            if byte != 0x00:
                length = length * 128 + (byte & 0x7f)
                if (byte & 0x80) != 0x80:
                    break

        return length

    def _read_until(self):
        bytes = []

        while True:
            byte = self.rfile.read(1)
            if ord(byte) != 0xff:
                bytes.append(byte)
            else:
                break

        return ''.join(bytes)

    def wait(self):
        while True:
            if self.websocket_closed:
                return None

            frame_str = self.rfile.read(1)
            if not frame_str:
                # Connection lost?
                self.websocket_closed = True
                continue
            else:
                frame_type = ord(frame_str)

            if (frame_type & 0x80) == 0x00: # most significant byte is not set

                if frame_type == 0x00:
                    bytes = self._read_until()
                    return bytes.decode("utf-8", "replace")
                else:
                    self.websocket_closed = True

            elif (frame_type & 0x80) == 0x80: # most significant byte is set
                # Read binary data (forward-compatibility)
                if frame_type != 0xff:
                    self.websocket_closed = True
                else:
                    length = self._message_length()
                    if length == 0:
                        self.websocket_closed = True
                    else:
                        self.rfile.read(length) # discard the bytes
            else:
                raise IOError("Reveiced an invalid message")

class WebSocketVersion7(WebSocket):
    FIN = int("10000000", 2)
    RSV = int("01110000", 2)
    OPCODE = int("00001111", 2)
    MASK = int("10000000", 2)
    PAYLOAD = int("01111111", 2)

    OPCODE_FRAG = 0x0
    OPCODE_TEXT = 0x1
    OPCODE_BINARY = 0x2
    OPCODE_CLOSE = 0x8
    OPCODE_PING = 0x9
    OPCODE_PONG = 0xA

    REASON_NORMAL = 1000
    REASON_GOING_AWAY = 1001
    REASON_PROTOCOL_ERROR = 1002
    REASON_UNSUPPORTED_DATA_TYPE = 1003
    REASON_TOO_LARGE = 1004

    LEN_16 = 126
    LEN_64 = 127

    def __init__(self, sock, rfile, environ):
        self.rfile = rfile
        self.socket = sock
        self.origin = environ.get('HTTP_SEC_WEBSOCKET_ORIGIN')
        self.protocol = environ.get('HTTP_SEC_WEBSOCKET_PROTOCOL', 'unknown')
        self.path = environ.get('PATH_INFO')
        self.websocket_closed = False
        self._fragments = []

    def _read_from_socket(self, count):
        return self.rfile.read(count)

    # TODO: replace all magic numbers with constants
    def wait(self):
        while True:
            payload = ''
            if self.websocket_closed:
                return None

            opcode_octet, length_octet = struct.unpack('!BB', self._read_from_socket(2))

            if self.RSV & opcode_octet:
                self.close(self.REASON_PROTOCOL_ERROR, 'Reserved bits cannot be set')
                return None

            opcode = opcode_octet & self.OPCODE
            is_final_frag = (self.FIN & opcode_octet) != 0

            if self._is_opcode_invalid(opcode):
                self.close(self.REASON_PROTOCOL_ERROR, 'Invalid opcode %x' % opcode)
                return None

            if not is_final_frag and self.OPCODE_CLOSE <= opcode <= self.OPCODE_PONG:
                self.close(self.REASON_PROTOCOL_ERROR, 'Control frames cannot be fragmented')
                return None

            if len(self._fragments) > 0 and not is_final_frag and opcode != self.OPCODE_FRAG:
                self.close(self.REASON_PROTOCOL_ERROR,
                        'Received new fragment frame with non-zero opcode')
                return None

            if len(self._fragments) > 0 and is_final_frag and (
                    self.OPCODE_TEXT <= opcode <= self.OPCODE_BINARY):
                self.close(self.REASON_PROTOCOL_ERROR,
                        'Received new unfragmented data frame during fragmented message')
                return None

            if not self.MASK & length_octet:
                self.close(self.REASON_PROTOCOL_ERROR, 'MASK must be set')
                return None

            length_code = length_octet & self.PAYLOAD

            if length_code >= self.LEN_16 and (self.OPCODE_CLOSE <= opcode <= self.OPCODE_PONG):
                self.close(self.REASON_PROTOCOL_ERROR,
                        'Control frame payload cannot be larger than 125 bytes')
                return None

            if length_code < self.LEN_16:
                length = length_code
            elif length_code == self.LEN_16:
                length = struct.unpack('!H', self._read_from_socket(2))[0]
            elif length_code == self.LEN_64:
                length = struct.unpack('!Q', self._read_from_socket(8))[0]
            else:
                raise Exception('Calculated invalid length')

            mask_octets = struct.unpack('!BBBB', self._read_from_socket(4))
            masked_payload = self._read_from_socket(length)

            payload = ''

            j = 0
            for c in masked_payload:
                # TODO: optimize me? http://www.skymind.com/~ocrow/python_string/
                payload += chr(ord(c) ^ mask_octets[j])
                j = (j + 1) % 4

            if opcode == self.OPCODE_TEXT:
                payload = payload.decode('utf-8')
            elif opcode == self.OPCODE_CLOSE:
                if length >= 2:
                    reason, message = struct.unpack('!H%ds' % (length - 2), payload)
                else:
                    reason = message = None

                self.close(self.REASON_NORMAL, '')
                return (reason, message)

            if opcode == self.OPCODE_PING:
                self.send(self.OPCODE_PONG, payload)
                return (self.OPCODE_PING, payload)
            elif opcode == self.OPCODE_PONG:
                return (self.OPCODE_PONG, payload)

            if is_final_frag:
                payload = ''.join(self._fragments) + payload
                self._fragments = []
                return payload
            else:
                self._fragments.append(payload)

    def _encode_text(self, s):
        if isinstance(s, unicode):
            return s.encode('utf-8')
        elif isinstance(s, str):
            return unicode(s).encode('utf-8')
        else:
            raise Exception('Invalid encoding')

    def _is_opcode_invalid(self, opcode):
        return opcode < self.OPCODE_FRAG or (opcode > self.OPCODE_BINARY and
                opcode < self.OPCODE_CLOSE) or opcode > self.OPCODE_PONG

    def send(self, opcode, message):
        if self.websocket_closed:
            raise Exception('Connection was terminated')

        if self._is_opcode_invalid(opcode):
            raise Exception('Invalid opcode %d' % opcode)

        if opcode == self.OPCODE_TEXT:
            message = self._encode_text(message)

        length = len(message)

        if opcode == self.OPCODE_TEXT:
            message = struct.pack('!%ds' % length, message)

        if length < self.LEN_16:
            preamble = struct.pack('!BB', self.FIN | opcode, length)
        elif length < 2 ** 16:
            preamble = struct.pack('!BBH', self.FIN | opcode, self.LEN_16, length)
        elif length < 2 ** 64:
            preamble = struct.pack('!BBQ', self.FIN | opcode, self.LEN_64, length)
        else:
            # this can't really happen, but for correctness sake...
            raise Exception('Message is too long')

        self.socket.sendall(preamble + message)

    def close(self, reason, message):
        message = self._encode_text(message)
        self.send(self.OPCODE_CLOSE, struct.pack('!H%ds' % len(message), reason, message))
        self.websocket_closed = True

        # based on gevent/pywsgi.py
        # see http://pypi.python.org/pypi/gevent#downloads
        if self.socket is not None:
            try:
                self.socket._sock.close()
                self.socket.close()
            except socket.error:
                pass
<|endoftext|>