<commit_before># -*- coding: utf-8 -*-
import logging
import csv

from django.views.generic import TemplateView
from django.http import HttpResponse

from haystack.query import SearchQuerySet, EmptySearchQuerySet

from dynamiq.utils import get_advanced_search_formset_class, FormsetQBuilder, ParsedStringQBuilder

from .forms import MEPSearchForm, MEPSearchAdvancedFormset, MEPSimpleSearchForm
from .shortcuts import TopRated, WorstRated

log = logging.getLogger(__name__)

class SearchView(TemplateView):

    template_name = 'search/search.html'
    list_template_name = "blocks/representative_list.html"

    def get_template_names(self):
        """
        Dispatch template according to the kind of request: ajax or normal.
        """
        if self.request.is_ajax():
            return [self.list_template_name]
        else:
            return [self.template_name]

    def get_context_data(self, **kwargs):
        query = None
        sort = MEPSearchAdvancedFormset.options_form_class.SORT_INITIAL
        limit = MEPSearchAdvancedFormset.options_form_class.LIMIT_INITIAL
        label = ""

        formset_class = get_advanced_search_formset_class(self.request.user, MEPSearchAdvancedFormset, MEPSearchForm)
        if "q" in self.request.GET:
            form = MEPSimpleSearchForm(self.request.GET)
            if form.is_valid():
                F = ParsedStringQBuilder(form.cleaned_data['q'], MEPSearchForm)
                query, label = F()
                formset = formset_class()
                _limit = form.cleaned_data.get("limit")
                if _limit is not None: limit = _limit
                sort = form.cleaned_data.get("sort") or sort
        else:
            form = MEPSimpleSearchForm()
            formset = formset_class(self.request.GET or None)
            formset.full_clean()
            if formset.is_valid():
                F = FormsetQBuilder(formset)
                query, label = F()
                sort = formset.options_form.cleaned_data.get("sort", sort)
                limit = formset.options_form.cleaned_data.get("limit", limit)

        if query:
            results = SearchQuerySet().filter(query)
            if sort:
                results = results.order_by(sort)
            if not limit:
                # When iterating over SearchQuerySet, haystack will fetch
                # results 10 by 10. This fetchs them all in one call:
                results = results[:]

            # we must find the average score for the search results
            average = sum([mep.total_score for mep in results if mep.total_score])/len(filter(lambda mep: mep.total_score, results))
        else:
            results = EmptySearchQuerySet()
            average = 0.0
        return {
            "dynamiq": {
                "results": results,
                "label": label,
                "formset": formset,
                "form": form,
                "average": average,
                "shortcuts": [
                    TopRated({"request": self.request}),
                    WorstRated({"request": self.request})
                ]
            },
            "list_template_name": self.list_template_name,
            "per_page": limit,
            "as_csv":  self.request.GET.get('as_csv', False)
        }

    def render_to_response(self, context, **response_kwargs):
        if self.request.GET.get('as_csv', False):
            return self.render_to_csv(context, **response_kwargs)
        return super(SearchView, self).render_to_response(context,
                                                          **response_kwargs)

    def render_to_csv(self, context, **response_kwargs):
        params = self.request.GET
        response = HttpResponse(mimetype='text/csv')
        name = self.request.path.strip('/').replace('/', '_')
        response['Content-Disposition'] = 'attachment; filename=%s.csv' % name

        header = [
            'name',
            'gender',
            'score',
            'emails',
            'country',
            'group',
            'bxl building id',
            'bxl building name',
            'bxl floor',
            'bxl office number',
            'bxl fax',
            'bxl phone1',
            'bxl phone2',
            'stg building id',
            'stg building name',
            'stg floor',
            'stg office number',
            'stg fax',
            'stg phone1',
            'stg phone2',
        ]

        meps = []

        if 'object' in context:
            obj = context['object']
            meps = getattr(obj, 'meps', [])

            if hasattr(meps, 'query'):
                # got a queryset
                meps = meps.select_related().distinct()
                if 'group' in params:
                    meps = meps.filter(groups__abbreviation=params['group'])
                    if 'country' in params:
                        meps = meps.filter(countries__name=params['country'])

        delim = ','
        if 'delim' in self.request.GET:
            delim = self.request.GET['delim'].encode()
        quote = '"'
        if 'quote' in self.request.GET:
            quote = self.request.GET['quote'].encode()

        writer = csv.writer(response, delimiter=delim,
                            quotechar=quote,
                            doublequote=True)
        writer.writerow(header)

        data = self.get_context_data(**response_kwargs)

        for result in data['dynamiq']['results']:
            mep = result.object
            row = [
                unicode(mep),
                mep.gender,
                int(mep.total_score) if mep.total_score else '',
                u' - '.join(mep.emails),
                mep.country.name,
                mep.group.abbreviation,
                mep.bxl_building.id if mep.bxl_building else '',
                mep.bxl_building.name if mep.bxl_building else '',
                mep.bxl_floor,
                mep.bxl_office_number,
                mep.bxl_fax,
                mep.bxl_phone1,
                mep.bxl_phone2,
                mep.stg_building.id if mep.stg_building else '',
                mep.stg_building.name if mep.stg_building else '',
                mep.stg_floor,
                mep.stg_office_number,
                mep.stg_fax,
                mep.stg_phone1,
                mep.stg_phone2,
            ]
            str_row = []
            for v in row:
                if isinstance(v, unicode):
                    v = v.encode('utf-8')
                str_row.append(v)

            mep_committees = {}
            # FIXME: this should use current_committees but it look like its broken
            # or no longer valid
            roles = mep.committeerole_set.all()
            for role in roles:
                abbr = role.committee.abbreviation
                mep_committees.setdefault(abbr, set([])).add(role.role)
            # for c in committees:
            #     str_row.append(' - '.join(sorted(mep_committees.get(c, []))))

            assert len(str_row) == len(header)
            writer.writerow(str_row)

        return response

class XhrSearchView(SearchView):

    template_name = "search/xhr.html"
<commit_msg>[fix] avoid to divide by zero<commit_after># -*- coding: utf-8 -*-
import logging
import csv

from django.views.generic import TemplateView
from django.http import HttpResponse

from haystack.query import SearchQuerySet, EmptySearchQuerySet

from dynamiq.utils import get_advanced_search_formset_class, FormsetQBuilder, ParsedStringQBuilder

from .forms import MEPSearchForm, MEPSearchAdvancedFormset, MEPSimpleSearchForm
from .shortcuts import TopRated, WorstRated

log = logging.getLogger(__name__)

class SearchView(TemplateView):

    template_name = 'search/search.html'
    list_template_name = "blocks/representative_list.html"

    def get_template_names(self):
        """
        Dispatch template according to the kind of request: ajax or normal.
        """
        if self.request.is_ajax():
            return [self.list_template_name]
        else:
            return [self.template_name]

    def get_context_data(self, **kwargs):
        query = None
        sort = MEPSearchAdvancedFormset.options_form_class.SORT_INITIAL
        limit = MEPSearchAdvancedFormset.options_form_class.LIMIT_INITIAL
        label = ""

        formset_class = get_advanced_search_formset_class(self.request.user, MEPSearchAdvancedFormset, MEPSearchForm)
        if "q" in self.request.GET:
            form = MEPSimpleSearchForm(self.request.GET)
            if form.is_valid():
                F = ParsedStringQBuilder(form.cleaned_data['q'], MEPSearchForm)
                query, label = F()
                formset = formset_class()
                _limit = form.cleaned_data.get("limit")
                if _limit is not None: limit = _limit
                sort = form.cleaned_data.get("sort") or sort
        else:
            form = MEPSimpleSearchForm()
            formset = formset_class(self.request.GET or None)
            formset.full_clean()
            if formset.is_valid():
                F = FormsetQBuilder(formset)
                query, label = F()
                sort = formset.options_form.cleaned_data.get("sort", sort)
                limit = formset.options_form.cleaned_data.get("limit", limit)

        if query:
            results = SearchQuerySet().filter(query)
            if sort:
                results = results.order_by(sort)
            if not limit:
                # When iterating over SearchQuerySet, haystack will fetch
                # results 10 by 10. This fetchs them all in one call:
                results = results[:]

            # we must find the average score for the search results
            if len(filter(lambda mep: mep.total_score, results)) != 0:
                average = sum([mep.total_score for mep in results if mep.total_score])/len(filter(lambda mep: mep.total_score, results))
            else:
                average = 0.0
        else:
            results = EmptySearchQuerySet()
            average = 0.0
        return {
            "dynamiq": {
                "results": results,
                "label": label,
                "formset": formset,
                "form": form,
                "average": average,
                "shortcuts": [
                    TopRated({"request": self.request}),
                    WorstRated({"request": self.request})
                ]
            },
            "list_template_name": self.list_template_name,
            "per_page": limit,
            "as_csv":  self.request.GET.get('as_csv', False)
        }

    def render_to_response(self, context, **response_kwargs):
        if self.request.GET.get('as_csv', False):
            return self.render_to_csv(context, **response_kwargs)
        return super(SearchView, self).render_to_response(context,
                                                          **response_kwargs)

    def render_to_csv(self, context, **response_kwargs):
        params = self.request.GET
        response = HttpResponse(mimetype='text/csv')
        name = self.request.path.strip('/').replace('/', '_')
        response['Content-Disposition'] = 'attachment; filename=%s.csv' % name

        header = [
            'name',
            'gender',
            'score',
            'emails',
            'country',
            'group',
            'bxl building id',
            'bxl building name',
            'bxl floor',
            'bxl office number',
            'bxl fax',
            'bxl phone1',
            'bxl phone2',
            'stg building id',
            'stg building name',
            'stg floor',
            'stg office number',
            'stg fax',
            'stg phone1',
            'stg phone2',
        ]

        meps = []

        if 'object' in context:
            obj = context['object']
            meps = getattr(obj, 'meps', [])

            if hasattr(meps, 'query'):
                # got a queryset
                meps = meps.select_related().distinct()
                if 'group' in params:
                    meps = meps.filter(groups__abbreviation=params['group'])
                    if 'country' in params:
                        meps = meps.filter(countries__name=params['country'])

        delim = ','
        if 'delim' in self.request.GET:
            delim = self.request.GET['delim'].encode()
        quote = '"'
        if 'quote' in self.request.GET:
            quote = self.request.GET['quote'].encode()

        writer = csv.writer(response, delimiter=delim,
                            quotechar=quote,
                            doublequote=True)
        writer.writerow(header)

        data = self.get_context_data(**response_kwargs)

        for result in data['dynamiq']['results']:
            mep = result.object
            row = [
                unicode(mep),
                mep.gender,
                int(mep.total_score) if mep.total_score else '',
                u' - '.join(mep.emails),
                mep.country.name,
                mep.group.abbreviation,
                mep.bxl_building.id if mep.bxl_building else '',
                mep.bxl_building.name if mep.bxl_building else '',
                mep.bxl_floor,
                mep.bxl_office_number,
                mep.bxl_fax,
                mep.bxl_phone1,
                mep.bxl_phone2,
                mep.stg_building.id if mep.stg_building else '',
                mep.stg_building.name if mep.stg_building else '',
                mep.stg_floor,
                mep.stg_office_number,
                mep.stg_fax,
                mep.stg_phone1,
                mep.stg_phone2,
            ]
            str_row = []
            for v in row:
                if isinstance(v, unicode):
                    v = v.encode('utf-8')
                str_row.append(v)

            mep_committees = {}
            # FIXME: this should use current_committees but it look like its broken
            # or no longer valid
            roles = mep.committeerole_set.all()
            for role in roles:
                abbr = role.committee.abbreviation
                mep_committees.setdefault(abbr, set([])).add(role.role)
            # for c in committees:
            #     str_row.append(' - '.join(sorted(mep_committees.get(c, []))))

            assert len(str_row) == len(header)
            writer.writerow(str_row)

        return response

class XhrSearchView(SearchView):

    template_name = "search/xhr.html"
<|endoftext|>


<commit_before>#!/usr/bin/env python2
# -*-coding:UTF-8 -*

import ConfigParser
from ConfigParser import RawConfigParser
import os

def main():

    configfile = os.path.join(os.environ['AIL_BIN'], 'packages/config.cfg')
    if not os.path.exists(configfile):
        raise Exception('Unable to find the configuration file. \
                        Did you set environment variables? \
                        Or activate the virtualenv.')
    configfileSample  = os.path.join(os.environ['AIL_BIN'], 'packages/config.cfg.sample')

    cfg = ConfigParser.ConfigParser()
    cfg.read(configfile)
    cfgSample = ConfigParser.ConfigParser()
    cfgSample.read(configfileSample)

    sections = RawConfigParser.sections(cfg)
    sectionsSample = RawConfigParser.sections(cfgSample)

    missingSection = []
    dicoMissingSection = {}
    missingItem = []
    dicoMissingItem = {}

    for sec in sectionsSample:
        if sec not in sections:
            missingSection += [sec]
            dicoMissingSection[sec] = RawConfigParser.items(cfgSample, sec)
        else:
            setSample = set(RawConfigParser.options(cfgSample, sec))
            setNormal = set(RawConfigParser.options(cfg, sec))
            if setSample != setNormal:
                missing_items = list(setSample.difference(setNormal))
                missingItem += [sec]
                list_items = []
                for i in missing_items:
                    list_items.append( (i, cfgSample.get(sec, i)) )
                dicoMissingItem[sec] = list_items

    if len(missingSection) == 0 and len(missingItem) == 0:
        print("Configuration up-to-date")
        return
    print("/!\\ Configuration not complete. Missing following configuration: /!\\")
    print("+--------------------------------------------------------------------+")
    for section in missingSection:
        print("["+section+"]")
        for item in dicoMissingSection[section]:
            print("  - "+item[0])
    for section in missingItem:
        print("["+section+"]")
        for item in dicoMissingItem[section]:
            print("  - "+item[0])
    print("+--------------------------------------------------------------------+")
    resp = raw_input("Do you want to auto fix it? [y/n] ")

    if resp != 'y':
        return
    else:
        for section in missingSection:
            cfg.add_section(section)
            for item, value in dicoMissingSection[section]:
                cfg.set(section, item, value)
        for section in missingItem:
            for item, value in dicoMissingItem[section]:
                cfg.set(section, item, value)

        with open(configfile, 'w') as f:
            cfg.write(f)

if __name__ == "__main__":
    main()

<commit_msg>Keep original section order while updating<commit_after>#!/usr/bin/env python2
# -*-coding:UTF-8 -*

import ConfigParser
from ConfigParser import ConfigParser as cfgP
import os
from collections import OrderedDict

def main():

    configfile = os.path.join(os.environ['AIL_BIN'], 'packages/config.cfg')
    if not os.path.exists(configfile):
        raise Exception('Unable to find the configuration file. \
                        Did you set environment variables? \
                        Or activate the virtualenv.')
    configfileSample  = os.path.join(os.environ['AIL_BIN'], 'packages/config.cfg.sample')

    cfg = ConfigParser.ConfigParser()
    cfg.read(configfile)
    cfgSample = ConfigParser.ConfigParser()
    cfgSample.read(configfileSample)

    sections = cfgP.sections(cfg)
    sectionsSample = cfgP.sections(cfgSample)

    missingSection = []
    dicoMissingSection = {}
    missingItem = []
    dicoMissingItem = {}

    for sec in sectionsSample:
        if sec not in sections:
            missingSection += [sec]
            dicoMissingSection[sec] = cfgP.items(cfgSample, sec)
        else:
            setSample = set(cfgP.options(cfgSample, sec))
            setNormal = set(cfgP.options(cfg, sec))
            if setSample != setNormal:
                missing_items = list(setSample.difference(setNormal))
                missingItem += [sec]
                list_items = []
                for i in missing_items:
                    list_items.append( (i, cfgSample.get(sec, i)) )
                dicoMissingItem[sec] = list_items

    if len(missingSection) == 0 and len(missingItem) == 0:
        print("Configuration up-to-date")
        return
    print("/!\\ Configuration not complete. Missing following configuration: /!\\")
    print("+--------------------------------------------------------------------+")
    for section in missingSection:
        print("["+section+"]")
        for item in dicoMissingSection[section]:
            print("  - "+item[0])
    for section in missingItem:
        print("["+section+"]")
        for item in dicoMissingItem[section]:
            print("  - "+item[0])
    print("+--------------------------------------------------------------------+")

    resp = raw_input("Do you want to auto fix it? [y/n] ")

    if resp != 'y':
        return
    else:
        #Do not keep item order in section. New items appened
        for section in missingItem:
            for item, value in dicoMissingItem[section]:
                cfg.set(section, item, value)

        #Keep sections order while updating the config file
        new_dico = add_items_to_correct_position(cfgSample._sections, cfg._sections, missingSection, dicoMissingSection)
        cfg._sections = new_dico

        with open(configfile, 'w') as f:
            cfg.write(f)

''' Return a new dico with the section ordered as the old configuration with the updated one added '''
def add_items_to_correct_position(sample_dico, old_dico, missingSection, dicoMissingSection):
    new_dico = OrderedDict()

    positions = {}
    for pos_i, sec in enumerate(sample_dico):
        if sec in missingSection:
            positions[pos_i] = sec

    for pos_i, sec in enumerate(old_dico):
        if pos_i in positions:
            missSection = positions[pos_i]
            new_dico[missSection] = sample_dico[missSection]

        new_dico[sec] = old_dico[sec]
    return new_dico

if __name__ == "__main__":
    main()

<|endoftext|>


<commit_before>#!/usr/bin/env python
'''
unit testing for the model classes.

'''
__author__ = 'M@Campbell'

import unittest
from flask import url_for
from ooiservices.app import create_app, db
from ooiservices.app.models import Array, InstrumentDeployment, PlatformDeployment, Stream, \
StreamParameter, User, OperatorEvent, OperatorEventType, Organization

'''
These tests are additional to the normal testing performed by coverage; each of
these tests are to validate model logic outside of db management.

'''
class ModelTestCase(unittest.TestCase):
    def setUp(self):
        self.app = create_app('TESTING_CONFIG')
        self.app_context = self.app.app_context()
        self.app_context.push()
        db.create_all()
        self.client = self.app.test_client(use_cookies=False)
        Organization.insert_org()

    def tearDown(self):
        db.session.remove()
        db.drop_all()
        self.app_context.pop()

    def test_array(self):
        #Test the json in the object
        array = Array()
        self.assertTrue(array.to_json() == {'id': None, 'array_code': None, \
        'array_name': None, 'description': None, 'display_name': None, \
        'geo_location': None})

    def test_platform_deployment(self):
        #Test the json in the object
        platform_deployment = PlatformDeployment()
        self.assertTrue(platform_deployment.to_json() == {'id': None, \
        'array_id': None, 'display_name': None, 'end_date': None, \
        'geo_location': None, 'reference_designator': None, 'start_date': None})

    def test_instrument_deployment(self):
        #Test the json in the object
        instrument_deployment = InstrumentDeployment()
        should_be = {
            'id' :None,
            'depth': None,
            'display_name' : None,
            'end_date' : None,
            'geo_location': None,
            'platform_deployment_id' : None,
            'reference_designator' : None,
            'start_date' : None
        }
        self.assertEquals(instrument_deployment.to_json() , should_be)

    def test_stream(self):
        #Test the json in the object
        stream = Stream()
        self.assertTrue(stream.to_json() == {'id': None, 'description': None, \
        'instrument_id': None, 'stream_name': None})

    def test_parameter(self):
        #Test the json in the object
        stream_param = StreamParameter()
        self.assertTrue(stream_param.to_json() == {'id': None, 'data_type': None, \
        'long_name': None, 'parameter_name': None, 'short_name': None, \
        'standard_name': None, 'units': None})

    def test_user(self):
        #Test the json in the object
        user = User()
        self.assertEquals(user.to_json(), {
            'email': None,
            'id': None,
            'user_id': None,
            'active':None,
            'first_name': None,
            'last_name' : None,
            'organization_id' : None,
            'phone_alternate' : None,
            'phone_primary' : None,
            'scopes' : [],
            'role' : None,
            'user_name': None})

    def test_operator_event_type(self):
        #Test the json in the object
        operator_event_type = OperatorEventType()
        self.assertTrue(operator_event_type.to_json() == {'id': None, 'type_name': None, 'type_description': None})

    def test_geometry(self):
        platform_deployment = PlatformDeployment()
        platform_deployment.reference_designator = 'TEST0000'
        platform_deployment.geo_location = 'POINT(-70 40)'
        db.session.add(platform_deployment)
        db.session.commit()
        pd = PlatformDeployment.query.filter(PlatformDeployment.reference_designator=='TEST0000').first()
        self.assertEquals(pd.geojson, {'coordinates': [-70, 40], 'type': 'Point'})
<commit_msg>Fixes test_models.py<commit_after>#!/usr/bin/env python
'''
unit testing for the model classes.

'''
__author__ = 'M@Campbell'

import unittest
from flask import url_for
from ooiservices.app import create_app, db
from ooiservices.app.models import Array, InstrumentDeployment, PlatformDeployment, Stream, \
StreamParameter, User, OperatorEvent, OperatorEventType, Organization

'''
These tests are additional to the normal testing performed by coverage; each of
these tests are to validate model logic outside of db management.

'''
class ModelTestCase(unittest.TestCase):
    def setUp(self):
        self.app = create_app('TESTING_CONFIG')
        self.app_context = self.app.app_context()
        self.app_context.push()
        db.create_all()
        self.client = self.app.test_client(use_cookies=False)
        Organization.insert_org()

    def tearDown(self):
        db.session.remove()
        db.drop_all()
        self.app_context.pop()

    def test_array(self):
        #Test the json in the object
        array = Array()
        self.assertTrue(array.to_json() == {'id': None, 'array_code': None, \
        'array_name': None, 'description': None, 'display_name': None, \
        'geo_location': None})

    def test_platform_deployment(self):
        #Test the json in the object
        platform_deployment = PlatformDeployment()
        self.assertTrue(platform_deployment.to_json() == {'id': None, \
        'array_id': None, 'display_name': None, 'end_date': None, \
        'geo_location': None, 'reference_designator': None, 'start_date': None})

    def test_instrument_deployment(self):
        #Test the json in the object
        instrument_deployment = InstrumentDeployment()
        should_be = {
            'id' :None,
            'depth': None,
            'display_name' : None,
            'end_date' : None,
            'geo_location': None,
            'platform_deployment_id' : None,
            'reference_designator' : None,
            'start_date' : None
        }
        self.assertEquals(instrument_deployment.to_json() , should_be)

    def test_stream(self):
        #Test the json in the object
        stream = Stream()
        self.assertTrue(stream.to_json() == {'id': None, 'description': None, \
        'instrument_id': None, 'stream_name': None})

    def test_parameter(self):
        #Test the json in the object
        stream_param = StreamParameter()
        self.assertTrue(stream_param.to_json() == {'id': None, 'data_type': None, \
        'long_name': None, 'parameter_name': None, 'short_name': None, \
        'standard_name': None, 'units': None})

    def test_user(self):
        #Test the json in the object
        user = User()
        self.assertEquals(user.to_json(), {
            'email': None,
            'id': None,
            'user_id': None,
            'active':None,
            'first_name': None,
            'last_name' : None,
            'organization_id' : None,
            'phone_alternate' : None,
            'phone_primary' : None,
            'scopes' : [],
            'role' : None,
            'user_name': None,
            'email_opt_in': None})

    def test_operator_event_type(self):
        #Test the json in the object
        operator_event_type = OperatorEventType()
        self.assertTrue(operator_event_type.to_json() == {'id': None, 'type_name': None, 'type_description': None})

    def test_geometry(self):
        platform_deployment = PlatformDeployment()
        platform_deployment.reference_designator = 'TEST0000'
        platform_deployment.geo_location = 'POINT(-70 40)'
        db.session.add(platform_deployment)
        db.session.commit()
        pd = PlatformDeployment.query.filter(PlatformDeployment.reference_designator=='TEST0000').first()
        self.assertEquals(pd.geojson, {'coordinates': [-70, 40], 'type': 'Point'})
<|endoftext|>


<commit_before>import os
from StringIO import StringIO
from tempfile import TemporaryFile, mkstemp

from twisted.trial import unittest
from twisted.internet import defer, reactor
from twisted.python.usage import UsageError

from ooni.nettest import NetTest, InvalidOption, MissingRequiredOption
from ooni.nettest import NetTestLoader, FailureToLoadNetTest
from ooni.tasks import BaseTask

from ooni.director import Director
from ooni.managers import TaskManager

from ooni.tests.mocks import MockMeasurement, MockMeasurementFailOnce
from ooni.tests.mocks import MockNetTest, MockDirector, MockReporter
from ooni.tests.mocks import MockMeasurementManager
defer.setDebugging(True)

net_test_string = """
from twisted.python import usage
from ooni.nettest import NetTestCase

class UsageOptions(usage.Options):
    optParameters = [['spam', 's', None, 'ham']]

class DummyTestCase(NetTestCase):

    usageOptions = UsageOptions

    def test_a(self):
        self.report['bar'] = 'bar'

    def test_b(self):
        self.report['foo'] = 'foo'
"""

net_test_root_required = net_test_string+"""
    requiresRoot = True
"""

net_test_string_with_file = """
from twisted.python import usage
from ooni.nettest import NetTestCase

class UsageOptions(usage.Options):
    optParameters = [['spam', 's', None, 'ham']]

class DummyTestCase(NetTestCase):
    inputFile = ['file', 'f', None, 'The input File']

    usageOptions = UsageOptions

    def test_a(self):
        self.report['bar'] = 'bar'

    def test_b(self):
        self.report['foo'] = 'foo'
"""

net_test_string_with_required_option = """
from twisted.python import usage
from ooni.nettest import NetTestCase

class UsageOptions(usage.Options):
    optParameters = [['spam', 's', None, 'ham'],
                     ['foo', 'o', None, 'moo'],
                     ['bar', 'o', None, 'baz'],
    ]

class DummyTestCase(NetTestCase):
    inputFile = ['file', 'f', None, 'The input File']

    usageOptions = UsageOptions

    def test_a(self):
        self.report['bar'] = 'bar'

    def test_b(self):
        self.report['foo'] = 'foo'

    requiredOptions = ['foo', 'bar']
"""

dummyInputs = range(1)
dummyArgs = ('--spam', 'notham')
dummyOptions = {'spam':'notham'}
dummyInvalidArgs = ('--cram', 'jam')
dummyInvalidOptions= {'cram':'jam'}
dummyArgsWithRequiredOptions = ('--foo', 'moo', '--bar', 'baz')
dummyRequiredOptions = {'foo':'moo', 'bar':'baz'}
dummyArgsWithFile = ('--spam', 'notham', '--file', 'dummyInputFile.txt')

class TestNetTest(unittest.TestCase):
    timeout = 1
    def setUp(self):
        with open('dummyInputFile.txt', 'w') as f:
            for i in range(10):
                f.write("%s\n" % i)

        from ooni.settings import config
        config.read_config_file()

    def assertCallable(self, thing):
        self.assertIn('__call__', dir(thing))

    def verifyMethods(self, testCases):
        uniq_test_methods = set()
        for test_class, test_methods in testCases:
            instance = test_class()
            for test_method in test_methods:
                c = getattr(instance, test_method)
                self.assertCallable(c)
                uniq_test_methods.add(test_method)
        self.assertEqual(set(['test_a', 'test_b']), uniq_test_methods)

    def test_load_net_test_from_file(self):
        """
        Given a file verify that the net test cases are properly
        generated.
        """
        __, net_test_file = mkstemp()
        with open(net_test_file, 'w') as f:
            f.write(net_test_string)
        f.close()

        ntl = NetTestLoader(dummyArgs)
        ntl.loadNetTestFile(net_test_file)

        self.verifyMethods(ntl.testCases)
        os.unlink(net_test_file)

    def test_load_net_test_from_str(self):
        """
        Given a file like object verify that the net test cases are properly
        generated.
        """
        ntl = NetTestLoader(dummyArgs)
        ntl.loadNetTestString(net_test_string)

        self.verifyMethods(ntl.testCases)

    def test_load_net_test_from_StringIO(self):
        """
        Given a file like object verify that the net test cases are properly
        generated.
        """
        ntl = NetTestLoader(dummyArgs)
        ntl.loadNetTestString(net_test_string)

        self.verifyMethods(ntl.testCases)

    def test_load_with_option(self):
        ntl = NetTestLoader(dummyArgs)
        ntl.loadNetTestString(net_test_string)

        self.assertIsInstance(ntl, NetTestLoader)
        for test_klass, test_meth in ntl.testCases:
            for option in dummyOptions.keys():
                self.assertIn(option, test_klass.usageOptions())

    def test_load_with_invalid_option(self):
        try:
            ntl = NetTestLoader(dummyInvalidArgs)
            ntl.loadNetTestString(net_test_string)

            ntl.checkOptions()
            raise Exception
        except UsageError:
            pass

    def test_load_with_required_option(self):
        ntl = NetTestLoader(dummyArgsWithRequiredOptions)
        ntl.loadNetTestString(net_test_string_with_required_option)

        self.assertIsInstance(ntl, NetTestLoader)

    def test_load_with_missing_required_option(self):
        try:
            ntl = NetTestLoader(dummyArgs)
            ntl.loadNetTestString(net_test_string_with_required_option)

        except MissingRequiredOption:
            pass

    def test_net_test_inputs(self):
        ntl = NetTestLoader(dummyArgsWithFile)
        ntl.loadNetTestString(net_test_string_with_file)

        ntl.checkOptions()

        # XXX: if you use the same test_class twice you will have consumed all
        # of its inputs!
        tested = set([])
        for test_class, test_method in ntl.testCases:
            if test_class not in tested:
                tested.update([test_class])
                self.assertEqual(len(list(test_class.inputs)), 10)

    def test_setup_local_options_in_test_cases(self):
        ntl = NetTestLoader(dummyArgs)
        ntl.loadNetTestString(net_test_string)

        ntl.checkOptions()

        for test_class, test_method in ntl.testCases:
            self.assertEqual(test_class.localOptions, dummyOptions)

    def test_generate_measurements_size(self):
        ntl = NetTestLoader(dummyArgsWithFile)
        ntl.loadNetTestString(net_test_string_with_file)

        ntl.checkOptions()
        net_test = NetTest(ntl, None)

        measurements = list(net_test.generateMeasurements())
        self.assertEqual(len(measurements), 20)

    def test_net_test_completed_callback(self):
        ntl = NetTestLoader(dummyArgsWithFile)
        ntl.loadNetTestString(net_test_string_with_file)

        ntl.checkOptions()
        director = Director()

        d = director.startNetTest(ntl, [MockReporter()])

        @d.addCallback
        def complete(result):
            self.assertEqual(result, None)
            self.assertEqual(director.successfulMeasurements, 20)

        return d

    def test_require_root_succeed(self):
        #XXX: will require root to run
        ntl = NetTestLoader(dummyArgs)
        ntl.loadNetTestString(net_test_root_required)

        for test_class, method in ntl.testCases:
            self.assertTrue(test_class.requiresRoot)
<commit_msg>Update Unit Tests<commit_after>import os
from StringIO import StringIO
from tempfile import TemporaryFile, mkstemp

from twisted.trial import unittest
from twisted.internet import defer, reactor
from twisted.python.usage import UsageError

from ooni.nettest import NetTest, InvalidOption, MissingRequiredOption
from ooni.nettest import NetTestLoader, FailureToLoadNetTest
from ooni.tasks import BaseTask

from ooni.director import Director
from ooni.managers import TaskManager

from ooni.tests.mocks import MockMeasurement, MockMeasurementFailOnce
from ooni.tests.mocks import MockNetTest, MockDirector, MockReporter
from ooni.tests.mocks import MockMeasurementManager
defer.setDebugging(True)

net_test_string = """
from twisted.python import usage
from ooni.nettest import NetTestCase

class UsageOptions(usage.Options):
    optParameters = [['spam', 's', None, 'ham']]

class DummyTestCase(NetTestCase):

    usageOptions = UsageOptions

    def test_a(self):
        self.report['bar'] = 'bar'

    def test_b(self):
        self.report['foo'] = 'foo'
"""

net_test_root_required = net_test_string+"""
    requiresRoot = True
"""

net_test_string_with_file = """
from twisted.python import usage
from ooni.nettest import NetTestCase

class UsageOptions(usage.Options):
    optParameters = [['spam', 's', None, 'ham']]

class DummyTestCase(NetTestCase):
    inputFile = ['file', 'f', None, 'The input File']

    usageOptions = UsageOptions

    def test_a(self):
        self.report['bar'] = 'bar'

    def test_b(self):
        self.report['foo'] = 'foo'
"""

net_test_string_with_required_option = """
from twisted.python import usage
from ooni.nettest import NetTestCase

class UsageOptions(usage.Options):
    optParameters = [['spam', 's', None, 'ham'],
                     ['foo', 'o', None, 'moo'],
                     ['bar', 'o', None, 'baz'],
    ]

class DummyTestCase(NetTestCase):
    inputFile = ['file', 'f', None, 'The input File']

    usageOptions = UsageOptions

    def test_a(self):
        self.report['bar'] = 'bar'

    def test_b(self):
        self.report['foo'] = 'foo'

    requiredOptions = ['foo', 'bar']
"""

dummyInputs = range(1)
dummyArgs = ('--spam', 'notham')
dummyOptions = {'spam':'notham'}
dummyInvalidArgs = ('--cram', 'jam')
dummyInvalidOptions= {'cram':'jam'}
dummyArgsWithRequiredOptions = ('--foo', 'moo', '--bar', 'baz')
dummyRequiredOptions = {'foo':'moo', 'bar':'baz'}
dummyArgsWithFile = ('--spam', 'notham', '--file', 'dummyInputFile.txt')

class TestNetTest(unittest.TestCase):
    timeout = 1
    def setUp(self):
        with open('dummyInputFile.txt', 'w') as f:
            for i in range(10):
                f.write("%s\n" % i)

        from ooni.settings import config
        config.read_config_file()

    def assertCallable(self, thing):
        self.assertIn('__call__', dir(thing))

    def verifyMethods(self, testCases):
        uniq_test_methods = set()
        for test_class, test_methods in testCases:
            instance = test_class()
            for test_method in test_methods:
                c = getattr(instance, test_method)
                self.assertCallable(c)
                uniq_test_methods.add(test_method)
        self.assertEqual(set(['test_a', 'test_b']), uniq_test_methods)

    def test_load_net_test_from_file(self):
        """
        Given a file verify that the net test cases are properly
        generated.
        """
        __, net_test_file = mkstemp()
        with open(net_test_file, 'w') as f:
            f.write(net_test_string)
        f.close()

        ntl = NetTestLoader(dummyArgs)
        ntl.loadNetTestFile(net_test_file)

        self.verifyMethods(ntl.testCases)
        os.unlink(net_test_file)

    def test_load_net_test_from_str(self):
        """
        Given a file like object verify that the net test cases are properly
        generated.
        """
        ntl = NetTestLoader(dummyArgs)
        ntl.loadNetTestString(net_test_string)

        self.verifyMethods(ntl.testCases)

    def test_load_net_test_from_StringIO(self):
        """
        Given a file like object verify that the net test cases are properly
        generated.
        """
        ntl = NetTestLoader(dummyArgs)
        ntl.loadNetTestString(net_test_string)

        self.verifyMethods(ntl.testCases)

    def test_load_with_option(self):
        ntl = NetTestLoader(dummyArgs)
        ntl.loadNetTestString(net_test_string)

        self.assertIsInstance(ntl, NetTestLoader)
        for test_klass, test_meth in ntl.testCases:
            for option in dummyOptions.keys():
                self.assertIn(option, test_klass.usageOptions())

    def test_load_with_invalid_option(self):
        try:
            ntl = NetTestLoader(dummyInvalidArgs)
            ntl.loadNetTestString(net_test_string)

            ntl.checkOptions()
            raise Exception
        except UsageError:
            pass

    def test_load_with_required_option(self):
        ntl = NetTestLoader(dummyArgsWithRequiredOptions)
        ntl.loadNetTestString(net_test_string_with_required_option)

        self.assertIsInstance(ntl, NetTestLoader)

    def test_load_with_missing_required_option(self):
        try:
            ntl = NetTestLoader(dummyArgs)
            ntl.loadNetTestString(net_test_string_with_required_option)

        except MissingRequiredOption:
            pass

    def test_net_test_inputs(self):
        ntl = NetTestLoader(dummyArgsWithFile)
        ntl.loadNetTestString(net_test_string_with_file)

        ntl.checkOptions()
        nt = NetTest(ntl,None)
        nt.initializeInputProcessor()

        # XXX: if you use the same test_class twice you will have consumed all
        # of its inputs!
        tested = set([])
        for test_class, test_method in ntl.testCases:
            if test_class not in tested:
                tested.update([test_class])
                self.assertEqual(len(list(test_class.inputs)), 10)

    def test_setup_local_options_in_test_cases(self):
        ntl = NetTestLoader(dummyArgs)
        ntl.loadNetTestString(net_test_string)

        ntl.checkOptions()

        for test_class, test_method in ntl.testCases:
            self.assertEqual(test_class.localOptions, dummyOptions)

    def test_generate_measurements_size(self):
        ntl = NetTestLoader(dummyArgsWithFile)
        ntl.loadNetTestString(net_test_string_with_file)

        ntl.checkOptions()
        net_test = NetTest(ntl, None)

        net_test.initializeInputProcessor()
        measurements = list(net_test.generateMeasurements())
        self.assertEqual(len(measurements), 20)

    def test_net_test_completed_callback(self):
        ntl = NetTestLoader(dummyArgsWithFile)
        ntl.loadNetTestString(net_test_string_with_file)

        ntl.checkOptions()
        director = Director()

        d = director.startNetTest(ntl, [MockReporter()])

        @d.addCallback
        def complete(result):
            self.assertEqual(result, None)
            self.assertEqual(director.successfulMeasurements, 20)

        return d

    def test_require_root_succeed(self):
        #XXX: will require root to run
        ntl = NetTestLoader(dummyArgs)
        ntl.loadNetTestString(net_test_root_required)

        for test_class, method in ntl.testCases:
            self.assertTrue(test_class.requiresRoot)
<|endoftext|>


<commit_before>from django.contrib.auth import authenticate, login
from django.contrib.auth.decorators import login_required
from django.http import HttpResponse
from django.shortcuts import redirect, render, get_object_or_404
from scheduler.models import *
import datetime
import parsedatetime.parsedatetime as pdt
from itertools import chain

def can_edit_post(user, post):
    if (user.has_perm('scheduler.add_edit_post') and
        (((user.has_perm('scheduler.approve_copy') or not post.pub_ready_copy) and
        (user.has_perm('scheduler.approve_online') or not post.pub_ready_online))
        or (user.has_perm('scheduler.approve_online')))
        and not post.sent):
        return True
    return False

@login_required
def user_settings(request):
    user = request.user
    message = {}

    if request.method == "POST":
        message = {
            "mtype":"success",
            "mtext":"Your information has been updated",
        }

        first_name = request.POST.get('first_name',None)
        last_name = request.POST.get('last_name',None)
        password1 = request.POST.get('password1',None)
        password2 = request.POST.get('password2',None)
        if first_name:
            user.first_name = first_name
        if last_name:
            user.last_name = last_name
        if password1 and password2 and (password1 == password2) and len(password1) > 5:
            user.set_password(password1)
        if password1 and password2 and (password1 != password2):
            message['mtext'] = "Your passwords don't match"
            message['mtype'] = "fail"

    context = {
        "user" : request.user,
        "message" : message,
    }
    return render(request, 'scheduler/user_settings.html', context)

@login_required
def dashboard(request):
    message = {}
    has_delete_permission = request.user.has_perm('scheduler.add_edit_post')
    if request.method == "POST" and has_delete_permission:
        post_id = request.POST.get('post_id_to_delete',None)
        post = get_object_or_404(SMPost, pk=post_id)
        post.delete()
        message = {
            "mtype":"success",
            "mtext":"Your post was deleted",
        }

    alt_date = None
    if request.method == "GET":
        date_change_str = request.GET.get('date',None)
        if date_change_str:
            cal = pdt.Calendar()
            date_change = cal.parse(date_change_str)
            if date_change[1] == 1 or date_change[1] == 3:
                alt_date = datetime.date(date_change[0][0], date_change[0][1], date_change[0][2])

    if alt_date:
        view_date = alt_date
    else:
        if datetime.datetime.now().hour <= 4:
            view_date = datetime.date.today()
        else:
            view_date = datetime.date.today() + datetime.timedelta(days=1)

    tomorrow_posts = SMPost.objects.filter(pub_date=view_date)
    lost_posts = SMPost.objects.filter(pub_date=None)

    context = {
        "user" : request.user,
        "sections" : Section.objects.all(),
        "smposts" : list(chain(tomorrow_posts, lost_posts)),
        "message" : message,
        "view_date" : view_date,
    }
    return render(request, 'scheduler/dashboard.html', context)

@login_required
def edit(request, post_id, post=None):
    if not post:
        post = get_object_or_404(SMPost, pk=post_id)

    if post.sent or not can_edit_post(request.user, post):
        message = {
            "mtype":"status",
        }

        if not can_edit_post(request.user, post):
            message['mtext'] = "You do not have permission to edit this post"

        if post.sent:
            message['mtext'] = "This post has already been sent and cannot be edited"

        context = {
            "user" : request.user,
            "sections" : Section.objects.all(),
            "post" : post,
            "message" : message,
        }
        return render(request, 'scheduler/view.html', context)

    message = {}
    if request.method == "POST" and can_edit_post(request.user, post):
        post.story_url = request.POST.get('url',None)
        post.slug = request.POST.get('slug',None)
        try:
            post.section = Section.objects.get(pk=request.POST.get('section',None))
        except:
            post.section = None
        post.post_twitter = request.POST.get('tweet',None)
        post.post_facebook = request.POST.get('fb',None)
        date_str = request.POST.get('pub_date',None)
        time_str = request.POST.get('pub_time',None)

        # Date
        cal = pdt.Calendar()
        date_parsed = cal.parse(date_str)
        if date_parsed[1] == 1 or date_parsed[1] == 3:
            post.pub_date = datetime.date(date_parsed[0][0], date_parsed[0][1], date_parsed[0][2])
        else:
            post.pub_date = None

        # Time
        time_parsed = cal.parse(time_str)
        if time_parsed[1] == 2 or time_parsed[1] == 3:
            post.pub_time = datetime.time(time_parsed[0][3], time_parsed[0][4])
        else:
            post.pub_time = None

        # Checkboxes
        if request.user.has_perm('scheduler.approve_copy'):
            if request.POST.get('approve-copy',False) == 'on':
                post.pub_ready_copy = True
            else:
                post.pub_ready_copy = False

        if request.user.has_perm('scheduler.approve_online'):
            if request.POST.get('approve-online',False) == 'on':
                post.pub_ready_online = True
            else:
                post.pub_ready_online = False

        post.save()
        message = {
            "mtype":"success",
            "mtext":"Your changes were saved!",
        }
    if request.method == "GET" and request.user.has_perm('scheduler.add_edit_post'):
        if request.GET.get('add',None) == "true":
            message = {
                "mtype":"success",
                "mtext":"Your post was successfully created!",
            }

    context = {
        "user" : request.user,
        "sections" : Section.objects.all(),
        "post" : post,
        "tomorrow" : datetime.date.today() + datetime.timedelta(days=1),
        "message" : message,
        "twitter_limit" : MeowSetting.objects.get(setting_key='twitter_character_limit').setting_value,
    }
    return render(request, 'scheduler/edit.html', context)

@login_required
def add(request):
    if not request.user.has_perm('scheduler.add_edit_post'):
        return redirect('/')
    post_id = -1
    if request.method == "POST" and request.user.has_perm('scheduler.add_edit_post'):
        post = SMPost()
        edit(request, -1, post)
        post_id = post.id
        return redirect("/edit/"+str(post.id)+"/?add=true")

    context = {
        "user" : request.user,
        "sections" : Section.objects.all(),
        "tomorrow" : datetime.date.today() + datetime.timedelta(days=1),
        "twitter_limit" : MeowSetting.objects.get(setting_key='twitter_character_limit').setting_value,
    }
    return render(request, 'scheduler/edit.html', context)

def logout(request):
    return logout(request)<commit_msg>Forgot to save the user data<commit_after>from django.contrib.auth import authenticate, login
from django.contrib.auth.decorators import login_required
from django.http import HttpResponse
from django.shortcuts import redirect, render, get_object_or_404
from scheduler.models import *
import datetime
import parsedatetime.parsedatetime as pdt
from itertools import chain

def can_edit_post(user, post):
    if (user.has_perm('scheduler.add_edit_post') and
        (((user.has_perm('scheduler.approve_copy') or not post.pub_ready_copy) and
        (user.has_perm('scheduler.approve_online') or not post.pub_ready_online))
        or (user.has_perm('scheduler.approve_online')))
        and not post.sent):
        return True
    return False

@login_required
def user_settings(request):
    user = request.user
    message = {}

    if request.method == "POST":
        message = {
            "mtype":"success",
            "mtext":"Your information has been updated",
        }

        first_name = request.POST.get('first_name',None)
        last_name = request.POST.get('last_name',None)
        password1 = request.POST.get('password1',None)
        password2 = request.POST.get('password2',None)
        if first_name:
            user.first_name = first_name
        if last_name:
            user.last_name = last_name
        if password1 and password2 and (password1 == password2) and len(password1) > 5:
            user.set_password(password1)
        if password1 and password2 and (password1 != password2):
            message['mtext'] = "Your passwords don't match"
            message['mtype'] = "fail"
        user.save()

    context = {
        "user" : request.user,
        "message" : message,
    }
    return render(request, 'scheduler/user_settings.html', context)

@login_required
def dashboard(request):
    message = {}
    has_delete_permission = request.user.has_perm('scheduler.add_edit_post')
    if request.method == "POST" and has_delete_permission:
        post_id = request.POST.get('post_id_to_delete',None)
        post = get_object_or_404(SMPost, pk=post_id)
        post.delete()
        message = {
            "mtype":"success",
            "mtext":"Your post was deleted",
        }

    alt_date = None
    if request.method == "GET":
        date_change_str = request.GET.get('date',None)
        if date_change_str:
            cal = pdt.Calendar()
            date_change = cal.parse(date_change_str)
            if date_change[1] == 1 or date_change[1] == 3:
                alt_date = datetime.date(date_change[0][0], date_change[0][1], date_change[0][2])

    if alt_date:
        view_date = alt_date
    else:
        if datetime.datetime.now().hour <= 4:
            view_date = datetime.date.today()
        else:
            view_date = datetime.date.today() + datetime.timedelta(days=1)

    tomorrow_posts = SMPost.objects.filter(pub_date=view_date)
    lost_posts = SMPost.objects.filter(pub_date=None)

    context = {
        "user" : request.user,
        "sections" : Section.objects.all(),
        "smposts" : list(chain(tomorrow_posts, lost_posts)),
        "message" : message,
        "view_date" : view_date,
    }
    return render(request, 'scheduler/dashboard.html', context)

@login_required
def edit(request, post_id, post=None):
    if not post:
        post = get_object_or_404(SMPost, pk=post_id)

    if post.sent or not can_edit_post(request.user, post):
        message = {
            "mtype":"status",
        }

        if not can_edit_post(request.user, post):
            message['mtext'] = "You do not have permission to edit this post"

        if post.sent:
            message['mtext'] = "This post has already been sent and cannot be edited"

        context = {
            "user" : request.user,
            "sections" : Section.objects.all(),
            "post" : post,
            "message" : message,
        }
        return render(request, 'scheduler/view.html', context)

    message = {}
    if request.method == "POST" and can_edit_post(request.user, post):
        post.story_url = request.POST.get('url',None)
        post.slug = request.POST.get('slug',None)
        try:
            post.section = Section.objects.get(pk=request.POST.get('section',None))
        except:
            post.section = None
        post.post_twitter = request.POST.get('tweet',None)
        post.post_facebook = request.POST.get('fb',None)
        date_str = request.POST.get('pub_date',None)
        time_str = request.POST.get('pub_time',None)

        # Date
        cal = pdt.Calendar()
        date_parsed = cal.parse(date_str)
        if date_parsed[1] == 1 or date_parsed[1] == 3:
            post.pub_date = datetime.date(date_parsed[0][0], date_parsed[0][1], date_parsed[0][2])
        else:
            post.pub_date = None

        # Time
        time_parsed = cal.parse(time_str)
        if time_parsed[1] == 2 or time_parsed[1] == 3:
            post.pub_time = datetime.time(time_parsed[0][3], time_parsed[0][4])
        else:
            post.pub_time = None

        # Checkboxes
        if request.user.has_perm('scheduler.approve_copy'):
            if request.POST.get('approve-copy',False) == 'on':
                post.pub_ready_copy = True
            else:
                post.pub_ready_copy = False

        if request.user.has_perm('scheduler.approve_online'):
            if request.POST.get('approve-online',False) == 'on':
                post.pub_ready_online = True
            else:
                post.pub_ready_online = False

        post.save()
        message = {
            "mtype":"success",
            "mtext":"Your changes were saved!",
        }
    if request.method == "GET" and request.user.has_perm('scheduler.add_edit_post'):
        if request.GET.get('add',None) == "true":
            message = {
                "mtype":"success",
                "mtext":"Your post was successfully created!",
            }

    context = {
        "user" : request.user,
        "sections" : Section.objects.all(),
        "post" : post,
        "tomorrow" : datetime.date.today() + datetime.timedelta(days=1),
        "message" : message,
        "twitter_limit" : MeowSetting.objects.get(setting_key='twitter_character_limit').setting_value,
    }
    return render(request, 'scheduler/edit.html', context)

@login_required
def add(request):
    if not request.user.has_perm('scheduler.add_edit_post'):
        return redirect('/')
    post_id = -1
    if request.method == "POST" and request.user.has_perm('scheduler.add_edit_post'):
        post = SMPost()
        edit(request, -1, post)
        post_id = post.id
        return redirect("/edit/"+str(post.id)+"/?add=true")

    context = {
        "user" : request.user,
        "sections" : Section.objects.all(),
        "tomorrow" : datetime.date.today() + datetime.timedelta(days=1),
        "twitter_limit" : MeowSetting.objects.get(setting_key='twitter_character_limit').setting_value,
    }
    return render(request, 'scheduler/edit.html', context)

def logout(request):
    return logout(request)<|endoftext|>


<commit_before>#!/usr/bin/env python
from optparse import OptionParser
import os
import subprocess
import sys
import tempfile

import h5py
import numpy as np
import pyBigWig
import tensorflow as tf

import basenji

from basenji_test import bigwig_open

'''
basenji_sed.py

Compute SNP expression difference scores for variants in a VCF file.

Note:
 -I'm having trouble verifying that I'm not double counting the scenario where
  two transcripts map to the same TSS. But I should verify that.
'''

################################################################################
# main
################################################################################
def main():
    usage = 'usage: %prog [options] <params_file> <model_file> <genes_hdf5_file> <vcf_file>'
    parser = OptionParser(usage)
    parser.add_option('-b', dest='batch_size', default=None, type='int', help='Batch size [Default: %default]')
    parser.add_option('-c', dest='csv', default=False, action='store_true', help='Print table as CSV [Default: %default]')
    parser.add_option('-g', dest='genome_file', default='%s/assembly/human.hg19.genome'%os.environ['HG19'], help='Chromosome lengths file [Default: %default]')
    parser.add_option('-i', dest='index_snp', default=False, action='store_true', help='SNPs are labeled with their index SNP as column 6 [Default: %default]')
    parser.add_option('-o', dest='out_dir', default='sed', help='Output directory for tables and plots [Default: %default]')
    parser.add_option('-s', dest='score', default=False, action='store_true', help='SNPs are labeled with scores as column 7 [Default: %default]')
    parser.add_option('-t', dest='target_wigs_file', default=None, help='Store target values, extracted from this list of WIG files')
    parser.add_option('--ti', dest='track_indexes', help='Comma-separated list of target indexes to output BigWig tracks')
    (options,args) = parser.parse_args()

    if len(args) != 4:
        parser.error('Must provide parameters and model files, genes HDF5 file, and QTL VCF file')
    else:
        params_file = args[0]
        model_file = args[1]
        genes_hdf5_file = args[2]
        vcf_file = args[3]

    if not os.path.isdir(options.out_dir):
        os.mkdir(options.out_dir)

    if options.track_indexes is None:
        options.track_indexes = []
    else:
        options.track_indexes = [int(ti) for ti in options.track_indexes.split(',')]
        if not os.path.isdir('%s/tracks' % options.out_dir):
            os.mkdir('%s/tracks' % options.out_dir)

    #################################################################
    # reads in genes HDF5

    genes_hdf5_in = h5py.File(genes_hdf5_file)

    seq_chrom = [chrom.decode('UTF-8') for chrom in genes_hdf5_in['seq_chrom']]
    seq_start = list(genes_hdf5_in['seq_start'])
    seq_end = list(genes_hdf5_in['seq_end'])
    seq_coords = list(zip(seq_chrom,seq_start,seq_end))

    seqs_1hot = genes_hdf5_in['seqs_1hot']
    print(seqs_1hot.shape)

    transcripts = [tx.decode('UTF-8') for tx in genes_hdf5_in['transcripts']]
    transcript_index = list(genes_hdf5_in['transcript_index'])
    transcript_pos = list(genes_hdf5_in['transcript_pos'])

    transcript_map = {}
    for ti in range(len(transcripts)):
        transcript_map[transcripts[ti]] = (transcript_index[ti], transcript_pos[ti])

    if 'transcript_targets' in genes_hdf5_in:
        transcript_targets = genes_hdf5_in['transcript_targets']
        target_labels = [tl.decode('UTF-8') for tl in genes_hdf5_in['target_labels']]
    else:
        transcript_targets = None
        target_labels = None

    #################################################################
    # prep SNPs

    # load SNPs
    snps = basenji.vcf.vcf_snps(vcf_file, options.index_snp, options.score, False)

    # intersect w/ segments
    snps_segs = intersect_snps(vcf_file, seq_coords)

    #################################################################
    # construct SNP sequences

    snp_seqs_1hot = []

    for snp_i in range(len(snps)):
        for seg_i in snps_segs[snp_i]:
            seq_chrom, seq_start, seq_end = seq_coords[seg_i]

            # determine the SNP's position in the segment
            snp_seq_pos = snps[snp_i].pos-1 - seq_start

            # write reference allele
            snp_seqs_1hot.append(seqs_1hot[seg_i])
            basenji.dna_io.hot1_set(snp_seqs_1hot[-1], snp_seq_pos, snps[snp_i].ref_allele)

            # write alternative allele
            snp_seqs_1hot.append(seqs_1hot[seg_i])
            basenji.dna_io.hot1_set(snp_seqs_1hot[-1], snp_seq_pos, snps[snp_i].alt_alleles[0])

    snp_seqs_1hot = np.array(snp_seqs_1hot)

    #################################################################
    # setup model

    job = basenji.dna_io.read_job_params(params_file)

    job['batch_length'] = snp_seqs_1hot.shape[1]
    job['seq_depth'] = snp_seqs_1hot.shape[2]
    job['target_pool'] = int(np.array(genes_hdf5_in['pool_width']))

    if transcript_targets is not None:
        job['num_targets'] = transcript_targets.shape[1]

    if 'num_targets' not in job:
        print("Must specify number of targets (num_targets) in the parameters file. I know, it's annoying. Sorry.", file=sys.stderr)
        exit(1)

    # build model
    model = basenji.rnn.RNN()
    model.build(job)

    if options.batch_size is not None:
        model.batch_size = options.batch_size

    # label targets
    if target_labels is None:
        if options.targets_file is None:
            target_labels = ['t%d' % ti for ti in range(job['num_targets'])]
        else:
            target_labels = [line.split()[0] for line in open(options.targets_file)]

    #################################################################
    # predict

    # initialize saver
    saver = tf.train.Saver()

    with tf.Session() as sess:
        # load variables into session
        saver.restore(sess, model_file)

        # initialize prediction stream
        seq_preds = Pred_Stream(sess, model, snp_seqs_1hot, 128)

        # determine prediction buffer
        pred_buffer = model.batch_buffer // model.target_pool

        #################################################################
        # collect and print SEDs

        header_cols = ('rsid', 'index', 'score', 'ref', 'alt', 'gene', 'target', 'ref_pred', 'alt pred', 'ser', 'sed')
        if options.csv:
            sed_out = open('%s/sed_table.csv' % options.out_dir, 'w')
            print(','.join(header_cols), file=sed_out)
        else:
            sed_out = open('%s/sed_table.txt' % options.out_dir, 'w')
            print(' '.join(header_cols), file=sed_out)

        pi = 0
        for snp_i in range(len(snps)):
            snp = snps[snp_i]

            # set index SNP
            snp_is = '%-13s' % '.'
            if options.index_snp:
                snp_is = '%-13s' % snp.index_snp

            # set score
            snp_score = '%5s' % '.'
            if options.score:
                snp_score = '%5.3f' % snp.score

            for seg_i in snps_segs[snp_i]:
                # get reference prediction (LxT)
                ref_preds = seq_preds.get(pi)
                pi += 1

                # get alternate prediction (LxT)
                alt_preds = seq_preds.get(pi)
                pi += 1

                # find genes
                for transcript in transcript_map:
                    tx_index, tx_pos = transcript_map[transcript]

                    if tx_index == seg_i:
                        # compute distance between SNP and gene
                        tx_gpos = seq_coords[1] + (tx_pos + 0.5)*model.target_pool
                        snp_dist = abs(tx_gpos - snp.pos)

                        # compute transcript pos in predictions
                        tx_pos_buf = tx_pos - pred_buffer

                        for ti in range(job['num_targets']):
                            # compute SED score
                            snp_gene_sed = (alt_preds[tx_pos_buf,ti] - ref_preds[tx_pos_buf,ti])
                            snp_gene_ser = np.divide(alt_preds[tx_pos_buf,ti]+1, ref_preds[tx_pos_buf,ti]+1)

                            # print to table
                            cols = (snp.rsid, snp_is, snp_score, basenji.vcf.cap_allele(snp.ref_allele), basenji.vcf.cap_allele(snp.alt_alleles[0]), transcript, target_labels[ti], ref_preds[tx_pos_buf,ti], alt_preds[tx_pos_buf,ti], snp_gene_ser, snp_gene_sed)
                            if options.csv:
                                print(','.join([str(c) for c in cols]), file=sed_out)
                            else:
                                print('%-13s %s %5s %6s %6s %12s %12s %6.4f %6.4f %7.4f' % cols, file=sed_out)

                # print tracks
                for ti in options.track_indexes:
                    ref_bw_file = '%s/tracks/%s_%s_t%d_ref.bw' % (options.out_dir, snp.rsid, seg_i, ti)
                    alt_bw_file = '%s/tracks/%s_%s_t%d_alt.bw' % (options.out_dir, snp.rsid, seg_i, ti)
                    ref_bw_open = bigwig_open(ref_bw_file, options.genome_file)
                    alt_bw_open = bigwig_open(alt_bw_file, options.genome_file)

                    seq_chrom, seq_start, seq_end = seq_coords[seg_i]
                    bw_chroms = [seq_chrom]*ref_preds.shape[0]
                    bw_starts = [int(seq_start + model.batch_buffer + bi*model.target_pool) for bi in range(ref_preds.shape[0])]
                    bw_ends = [int(bws + model.target_pool) for bws in bw_starts]

                    ref_values = [float(p) for p in ref_preds[:,ti]]
                    # print(ref_bw_open.chroms())
                    # print('bw_chroms', len(bw_chroms), type(bw_chroms[0]), bw_chroms[:10])
                    # print('bw_starts', len(bw_starts), type(bw_starts[0]), bw_starts[:10])
                    # print('bw_ends', len(bw_ends), type(bw_ends[0]), bw_ends[:10])
                    # print('ref_values', len(ref_values), type(ref_values[0]), ref_values[:10])
                    ref_bw_open.addEntries(bw_chroms, bw_starts, ends=bw_ends, values=ref_values)

                    alt_values = [float(p) for p in alt_preds[:,ti]]
                    alt_bw_open.addEntries(bw_chroms, bw_starts, ends=bw_ends, values=alt_values)

                    ref_bw_open.close()
                    alt_bw_open.close()

    sed_out.close()

    #################################################################
    # clean up

    genes_hdf5_in.close()

def intersect_snps(vcf_file, seq_coords):
    # print segments to BED
    # hash segments to indexes
    seg_temp = tempfile.NamedTemporaryFile()
    seg_bed_file = seg_temp.name
    seg_bed_out = open(seg_bed_file, 'w')
    segment_indexes = {}
    for si in range(len(seq_coords)):
        segment_indexes[seq_coords[si]] = si
        print('%s\t%d\t%d' % seq_coords[si], file=seg_bed_out)
    seg_bed_out.close()

    # hash SNPs to indexes
    snp_indexes = {}
    si = 0

    vcf_in = open(vcf_file)
    line = vcf_in.readline()
    while line[0] == '#':
        line = vcf_in.readline()
    while line:
        a = line.split()
        snp_id = a[2]
        if snp_id in snp_indexes:
            print('Duplicate SNP id %s will break the script' % snp_id, file=sys.stderr)
            exit(1)
        snp_indexes[snp_id] = si
        si += 1
        line = vcf_in.readline()
    vcf_in.close()

    # initialize list of lists
    snp_segs = []
    for i in range(len(snp_indexes)):
        snp_segs.append([])

    # intersect
    p = subprocess.Popen('bedtools intersect -wo -a %s -b %s' % (vcf_file, seg_bed_file), shell=True, stdout=subprocess.PIPE)
    for line in p.stdout:
        line = line.decode('UTF-8')
        a = line.split()
        snp_id = a[2]
        seg_chrom = a[-4]
        seg_start = int(a[-3])
        seg_end = int(a[-2])
        seg_key = (seg_chrom,seg_start,seg_end)

        snp_segs[snp_indexes[snp_id]].append(segment_indexes[seg_key])

    p.communicate()

    return snp_segs

class Pred_Stream:
    ''' Interface to acquire predictions via a buffered stream mechanism
         rather than getting them all at once and using excessive memory. '''

    def __init__(self, sess, model, seqs_1hot, stream_length):
        self.sess = sess
        self.model = model

        self.seqs_1hot = seqs_1hot

        self.stream_length = stream_length
        self.stream_start = 0
        self.stream_end = 0

        if self.stream_length % self.model.batch_size != 0:
            print('Make the stream length a multiple of the batch size', file=sys.stderr)
            exit(1)

    def get(self, i):
        # acquire predictions, if needed
        if i >= self.stream_end:
            self.stream_start = self.stream_end
            self.stream_end = min(self.stream_start + self.stream_length, self.seqs_1hot.shape[0])

            # subset sequences
            stream_seqs_1hot = self.seqs_1hot[self.stream_start:self.stream_end]

            # initialize batcher
            batcher = basenji.batcher.Batcher(stream_seqs_1hot, batch_size=self.model.batch_size)

            # predict
            self.stream_preds = self.model.predict(self.sess, batcher)
            # print('Acquired predictions ', self.stream_preds.shape, flush=True)

        return self.stream_preds[i - self.stream_start]

################################################################################
# __main__
################################################################################
if __name__ == '__main__':
    main()
<commit_msg>print more into table<commit_after>#!/usr/bin/env python
from optparse import OptionParser
import os
import subprocess
import sys
import tempfile

import h5py
import numpy as np
import pyBigWig
import tensorflow as tf

import basenji

from basenji_test import bigwig_open

'''
basenji_sed.py

Compute SNP expression difference scores for variants in a VCF file.

Note:
 -I'm having trouble verifying that I'm not double counting the scenario where
  two transcripts map to the same TSS. But I should verify that.
'''

################################################################################
# main
################################################################################
def main():
    usage = 'usage: %prog [options] <params_file> <model_file> <genes_hdf5_file> <vcf_file>'
    parser = OptionParser(usage)
    parser.add_option('-b', dest='batch_size', default=None, type='int', help='Batch size [Default: %default]')
    parser.add_option('-c', dest='csv', default=False, action='store_true', help='Print table as CSV [Default: %default]')
    parser.add_option('-g', dest='genome_file', default='%s/assembly/human.hg19.genome'%os.environ['HG19'], help='Chromosome lengths file [Default: %default]')
    parser.add_option('-i', dest='index_snp', default=False, action='store_true', help='SNPs are labeled with their index SNP as column 6 [Default: %default]')
    parser.add_option('-o', dest='out_dir', default='sed', help='Output directory for tables and plots [Default: %default]')
    parser.add_option('-s', dest='score', default=False, action='store_true', help='SNPs are labeled with scores as column 7 [Default: %default]')
    parser.add_option('-t', dest='target_wigs_file', default=None, help='Store target values, extracted from this list of WIG files')
    parser.add_option('--ti', dest='track_indexes', help='Comma-separated list of target indexes to output BigWig tracks')
    (options,args) = parser.parse_args()

    if len(args) != 4:
        parser.error('Must provide parameters and model files, genes HDF5 file, and QTL VCF file')
    else:
        params_file = args[0]
        model_file = args[1]
        genes_hdf5_file = args[2]
        vcf_file = args[3]

    if not os.path.isdir(options.out_dir):
        os.mkdir(options.out_dir)

    if options.track_indexes is None:
        options.track_indexes = []
    else:
        options.track_indexes = [int(ti) for ti in options.track_indexes.split(',')]
        if not os.path.isdir('%s/tracks' % options.out_dir):
            os.mkdir('%s/tracks' % options.out_dir)

    #################################################################
    # reads in genes HDF5

    genes_hdf5_in = h5py.File(genes_hdf5_file)

    seq_chrom = [chrom.decode('UTF-8') for chrom in genes_hdf5_in['seq_chrom']]
    seq_start = list(genes_hdf5_in['seq_start'])
    seq_end = list(genes_hdf5_in['seq_end'])
    seq_coords = list(zip(seq_chrom,seq_start,seq_end))

    seqs_1hot = genes_hdf5_in['seqs_1hot']
    print(seqs_1hot.shape)

    transcripts = [tx.decode('UTF-8') for tx in genes_hdf5_in['transcripts']]
    transcript_index = list(genes_hdf5_in['transcript_index'])
    transcript_pos = list(genes_hdf5_in['transcript_pos'])

    transcript_map = {}
    for ti in range(len(transcripts)):
        transcript_map[transcripts[ti]] = (transcript_index[ti], transcript_pos[ti])

    if 'transcript_targets' in genes_hdf5_in:
        transcript_targets = genes_hdf5_in['transcript_targets']
        target_labels = [tl.decode('UTF-8') for tl in genes_hdf5_in['target_labels']]
    else:
        transcript_targets = None
        target_labels = None

    #################################################################
    # prep SNPs

    # load SNPs
    snps = basenji.vcf.vcf_snps(vcf_file, options.index_snp, options.score, False)

    # intersect w/ segments
    snps_segs = intersect_snps(vcf_file, seq_coords)

    #################################################################
    # construct SNP sequences

    snp_seqs_1hot = []

    for snp_i in range(len(snps)):
        for seg_i in snps_segs[snp_i]:
            seq_chrom, seq_start, seq_end = seq_coords[seg_i]

            # determine the SNP's position in the segment
            snp_seq_pos = snps[snp_i].pos-1 - seq_start

            # write reference allele
            snp_seqs_1hot.append(seqs_1hot[seg_i])
            basenji.dna_io.hot1_set(snp_seqs_1hot[-1], snp_seq_pos, snps[snp_i].ref_allele)

            # write alternative allele
            snp_seqs_1hot.append(seqs_1hot[seg_i])
            basenji.dna_io.hot1_set(snp_seqs_1hot[-1], snp_seq_pos, snps[snp_i].alt_alleles[0])

    snp_seqs_1hot = np.array(snp_seqs_1hot)

    #################################################################
    # setup model

    job = basenji.dna_io.read_job_params(params_file)

    job['batch_length'] = snp_seqs_1hot.shape[1]
    job['seq_depth'] = snp_seqs_1hot.shape[2]
    job['target_pool'] = int(np.array(genes_hdf5_in['pool_width']))

    if transcript_targets is not None:
        job['num_targets'] = transcript_targets.shape[1]

    if 'num_targets' not in job:
        print("Must specify number of targets (num_targets) in the parameters file. I know, it's annoying. Sorry.", file=sys.stderr)
        exit(1)

    # build model
    model = basenji.rnn.RNN()
    model.build(job)

    if options.batch_size is not None:
        model.batch_size = options.batch_size

    # label targets
    if target_labels is None:
        if options.targets_file is None:
            target_labels = ['t%d' % ti for ti in range(job['num_targets'])]
        else:
            target_labels = [line.split()[0] for line in open(options.targets_file)]

    #################################################################
    # predict

    # initialize saver
    saver = tf.train.Saver()

    with tf.Session() as sess:
        # load variables into session
        saver.restore(sess, model_file)

        # initialize prediction stream
        seq_preds = Pred_Stream(sess, model, snp_seqs_1hot, 128)

        # determine prediction buffer
        pred_buffer = model.batch_buffer // model.target_pool

        #################################################################
        # collect and print SEDs

        header_cols = ('rsid', 'index', 'score', 'ref', 'alt', 'gene', 'target', 'ref_pred', 'alt pred', 'ser', 'sed')
        if options.csv:
            sed_out = open('%s/sed_table.csv' % options.out_dir, 'w')
            print(','.join(header_cols), file=sed_out)
        else:
            sed_out = open('%s/sed_table.txt' % options.out_dir, 'w')
            print(' '.join(header_cols), file=sed_out)

        pi = 0
        for snp_i in range(len(snps)):
            snp = snps[snp_i]

            # set index SNP
            snp_is = '%-13s' % '.'
            if options.index_snp:
                snp_is = '%-13s' % snp.index_snp

            # set score
            snp_score = '%5s' % '.'
            if options.score:
                snp_score = '%5.3f' % snp.score

            for seg_i in snps_segs[snp_i]:
                # get reference prediction (LxT)
                ref_preds = seq_preds.get(pi)
                pi += 1

                # get alternate prediction (LxT)
                alt_preds = seq_preds.get(pi)
                pi += 1

                # find genes
                for transcript in transcript_map:
                    tx_index, tx_pos = transcript_map[transcript]

                    if tx_index == seg_i:
                        # compute distance between SNP and gene
                        tx_gpos = seq_coords[seg_i][1] + (tx_pos + 0.5)*model.target_pool
                        snp_dist = abs(tx_gpos - snp.pos)

                        # compute transcript pos in predictions
                        tx_pos_buf = tx_pos - pred_buffer

                        for ti in range(job['num_targets']):
                            # compute SED score
                            snp_gene_sed = (alt_preds[tx_pos_buf,ti] - ref_preds[tx_pos_buf,ti])
                            snp_gene_ser = np.log2(alt_preds[tx_pos_buf,ti]+1) - np.log2(ref_preds[tx_pos_buf,ti]+1)

                            # print to table
                            cols = (snp.rsid, snp_is, snp_score, basenji.vcf.cap_allele(snp.ref_allele), basenji.vcf.cap_allele(snp.alt_alleles[0]), transcript, snp_dist, target_labels[ti], ref_preds[tx_pos_buf,ti], alt_preds[tx_pos_buf,ti], snp_gene_ser, snp_gene_sed)
                            if options.csv:
                                print(','.join([str(c) for c in cols]), file=sed_out)
                            else:
                                print('%-13s %s %5s %6s %6s %12s %5d %12s %6.4f %6.4f %7.4f %7.4f' % cols, file=sed_out)

                # print tracks
                for ti in options.track_indexes:
                    ref_bw_file = '%s/tracks/%s_%s_t%d_ref.bw' % (options.out_dir, snp.rsid, seg_i, ti)
                    alt_bw_file = '%s/tracks/%s_%s_t%d_alt.bw' % (options.out_dir, snp.rsid, seg_i, ti)
                    ref_bw_open = bigwig_open(ref_bw_file, options.genome_file)
                    alt_bw_open = bigwig_open(alt_bw_file, options.genome_file)

                    seq_chrom, seq_start, seq_end = seq_coords[seg_i]
                    bw_chroms = [seq_chrom]*ref_preds.shape[0]
                    bw_starts = [int(seq_start + model.batch_buffer + bi*model.target_pool) for bi in range(ref_preds.shape[0])]
                    bw_ends = [int(bws + model.target_pool) for bws in bw_starts]

                    ref_values = [float(p) for p in ref_preds[:,ti]]
                    # print(ref_bw_open.chroms())
                    # print('bw_chroms', len(bw_chroms), type(bw_chroms[0]), bw_chroms[:10])
                    # print('bw_starts', len(bw_starts), type(bw_starts[0]), bw_starts[:10])
                    # print('bw_ends', len(bw_ends), type(bw_ends[0]), bw_ends[:10])
                    # print('ref_values', len(ref_values), type(ref_values[0]), ref_values[:10])
                    ref_bw_open.addEntries(bw_chroms, bw_starts, ends=bw_ends, values=ref_values)

                    alt_values = [float(p) for p in alt_preds[:,ti]]
                    alt_bw_open.addEntries(bw_chroms, bw_starts, ends=bw_ends, values=alt_values)

                    ref_bw_open.close()
                    alt_bw_open.close()

    sed_out.close()

    #################################################################
    # clean up

    genes_hdf5_in.close()

def intersect_snps(vcf_file, seq_coords):
    # print segments to BED
    # hash segments to indexes
    seg_temp = tempfile.NamedTemporaryFile()
    seg_bed_file = seg_temp.name
    seg_bed_out = open(seg_bed_file, 'w')
    segment_indexes = {}
    for si in range(len(seq_coords)):
        segment_indexes[seq_coords[si]] = si
        print('%s\t%d\t%d' % seq_coords[si], file=seg_bed_out)
    seg_bed_out.close()

    # hash SNPs to indexes
    snp_indexes = {}
    si = 0

    vcf_in = open(vcf_file)
    line = vcf_in.readline()
    while line[0] == '#':
        line = vcf_in.readline()
    while line:
        a = line.split()
        snp_id = a[2]
        if snp_id in snp_indexes:
            print('Duplicate SNP id %s will break the script' % snp_id, file=sys.stderr)
            exit(1)
        snp_indexes[snp_id] = si
        si += 1
        line = vcf_in.readline()
    vcf_in.close()

    # initialize list of lists
    snp_segs = []
    for i in range(len(snp_indexes)):
        snp_segs.append([])

    # intersect
    p = subprocess.Popen('bedtools intersect -wo -a %s -b %s' % (vcf_file, seg_bed_file), shell=True, stdout=subprocess.PIPE)
    for line in p.stdout:
        line = line.decode('UTF-8')
        a = line.split()
        snp_id = a[2]
        seg_chrom = a[-4]
        seg_start = int(a[-3])
        seg_end = int(a[-2])
        seg_key = (seg_chrom,seg_start,seg_end)

        snp_segs[snp_indexes[snp_id]].append(segment_indexes[seg_key])

    p.communicate()

    return snp_segs

class Pred_Stream:
    ''' Interface to acquire predictions via a buffered stream mechanism
         rather than getting them all at once and using excessive memory. '''

    def __init__(self, sess, model, seqs_1hot, stream_length):
        self.sess = sess
        self.model = model

        self.seqs_1hot = seqs_1hot

        self.stream_length = stream_length
        self.stream_start = 0
        self.stream_end = 0

        if self.stream_length % self.model.batch_size != 0:
            print('Make the stream length a multiple of the batch size', file=sys.stderr)
            exit(1)

    def get(self, i):
        # acquire predictions, if needed
        if i >= self.stream_end:
            self.stream_start = self.stream_end
            self.stream_end = min(self.stream_start + self.stream_length, self.seqs_1hot.shape[0])

            # subset sequences
            stream_seqs_1hot = self.seqs_1hot[self.stream_start:self.stream_end]

            # initialize batcher
            batcher = basenji.batcher.Batcher(stream_seqs_1hot, batch_size=self.model.batch_size)

            # predict
            self.stream_preds = self.model.predict(self.sess, batcher)
            # print('Acquired predictions ', self.stream_preds.shape, flush=True)

        return self.stream_preds[i - self.stream_start]

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