<commit_before># Authors: Alexandre Gramfort <alexandre.gramfort@telecom-paristech.fr>
#          Denis Engemann <denis.engemann@gmail.com>
#          Martin Luessi <mluessi@nmr.mgh.harvard.edu>
#          Eric Larson <larson.eric.d@gmail.com>
#          Cathy Nangini <cnangini@gmail.com>
#          Mainak Jas <mainak@neuro.hut.fi>
#          Jona Sassenhagen <jona.sassenhagen@gmail.com>
#
# License: Simplified BSD

import os.path as op
import warnings

import numpy as np
from numpy.testing import assert_raises

from mne import io, read_events, Epochs, pick_types, read_cov
from mne.channels import read_layout
from mne.utils import slow_test, run_tests_if_main
from mne.viz.evoked import _butterfly_onselect, plot_compare_evokeds
from mne.viz.utils import _fake_click

# Set our plotters to test mode
import matplotlib
matplotlib.use('Agg')  # for testing don't use X server

warnings.simplefilter('always')  # enable b/c these tests throw warnings

base_dir = op.join(op.dirname(__file__), '..', '..', 'io', 'tests', 'data')
evoked_fname = op.join(base_dir, 'test-ave.fif')
raw_fname = op.join(base_dir, 'test_raw.fif')
cov_fname = op.join(base_dir, 'test-cov.fif')
event_name = op.join(base_dir, 'test-eve.fif')
event_id, tmin, tmax = 1, -0.1, 0.1
n_chan = 6
layout = read_layout('Vectorview-all')

def _get_raw():
    return io.read_raw_fif(raw_fname, preload=False)

def _get_events():
    return read_events(event_name)

def _get_picks(raw):
    return pick_types(raw.info, meg=True, eeg=False, stim=False,
                      ecg=False, eog=False, exclude='bads')

def _get_epochs():
    raw = _get_raw()
    raw.add_proj([], remove_existing=True)
    events = _get_events()
    picks = _get_picks(raw)
    # Use a subset of channels for plotting speed
    picks = picks[np.round(np.linspace(0, len(picks) - 1, n_chan)).astype(int)]
    picks[0] = 2  # make sure we have a magnetometer
    with warnings.catch_warnings(record=True):  # proj
        epochs = Epochs(raw, events[:5], event_id, tmin, tmax, picks=picks,
                        baseline=(None, 0))
    epochs.info['bads'] = [epochs.ch_names[-1]]
    return epochs

def _get_epochs_delayed_ssp():
    raw = _get_raw()
    events = _get_events()
    picks = _get_picks(raw)
    reject = dict(mag=4e-12)
    epochs_delayed_ssp = Epochs(raw, events[:10], event_id, tmin, tmax,
                                picks=picks, baseline=(None, 0),
                                proj='delayed', reject=reject)
    return epochs_delayed_ssp

@slow_test
def test_plot_evoked():
    """Test plotting of evoked
    """
    import matplotlib.pyplot as plt
    evoked = _get_epochs().average()
    with warnings.catch_warnings(record=True):
        fig = evoked.plot(proj=True, hline=[1], exclude=[], window_title='foo')
        # Test a click
        ax = fig.get_axes()[0]
        line = ax.lines[0]
        _fake_click(fig, ax,
                    [line.get_xdata()[0], line.get_ydata()[0]], 'data')
        _fake_click(fig, ax,
                    [ax.get_xlim()[0], ax.get_ylim()[1]], 'data')
        # plot with bad channels excluded & spatial_colors & zorder
        evoked.plot(exclude='bads')
        evoked.plot(exclude=evoked.info['bads'], spatial_colors=True, gfp=True,
                    zorder='std')

        # test selective updating of dict keys is working.
        evoked.plot(hline=[1], units=dict(mag='femto foo'))
        evoked_delayed_ssp = _get_epochs_delayed_ssp().average()
        evoked_delayed_ssp.plot(proj='interactive')
        evoked_delayed_ssp.apply_proj()
        assert_raises(RuntimeError, evoked_delayed_ssp.plot,
                      proj='interactive')
        evoked_delayed_ssp.info['projs'] = []
        assert_raises(RuntimeError, evoked_delayed_ssp.plot,
                      proj='interactive')
        assert_raises(RuntimeError, evoked_delayed_ssp.plot,
                      proj='interactive', axes='foo')
        plt.close('all')

        # test GFP only
        evoked.plot(gfp='only')
        assert_raises(ValueError, evoked.plot, gfp='foo')

        evoked.plot_image(proj=True)
        # plot with bad channels excluded
        evoked.plot_image(exclude='bads', cmap='interactive')
        evoked.plot_image(exclude=evoked.info['bads'])  # does the same thing
        plt.close('all')

        evoked.plot_topo()  # should auto-find layout
        _butterfly_onselect(0, 200, ['mag', 'grad'], evoked)
        plt.close('all')

        cov = read_cov(cov_fname)
        cov['method'] = 'empirical'
        evoked.plot_white(cov)
        evoked.plot_white([cov, cov])

        # plot_compare_evokeds: test condition contrast, CI, color assignment
        plot_compare_evokeds(evoked.copy().pick_types(meg='mag'))
        evoked.rename_channels({'MEG 2142': "MEG 1642"})
        # test picking & plotting grads
        plot_compare_evokeds(evoked, picks=3)

        colors = dict(red='r', blue='b')
        linestyles = dict(red='--', blue='-')
        red, blue = evoked.copy(), evoked.copy()
        red.data *= 1.1
        blue.data *= 0.9
        contrast = dict()
        contrast["red/stim"] = list((evoked.copy(), red))
        contrast["blue/stim"] = list((evoked.copy(), blue))
        # test a bunch of params at once
        plot_compare_evokeds(contrast, colors=colors, linestyles=linestyles,
                             picks=[0, 2], vlines=[.01, -.04], invert_y=True,
                             truncate_yaxis=False, ymax=10, ymin=-10,
                             styles={"red/stim": {"linewidth": 1}})
        assert_raises(ValueError, plot_compare_evokeds,
                      contrast, picks=[0, 3])  # bad picks: multiple types
        assert_raises(ValueError, plot_compare_evokeds,
                      evoked, colors=dict(fake=1))  # 'fake' not in conds
        assert_raises(ValueError, plot_compare_evokeds,
                      evoked, styles=dict(fake=1))  # 'fake' not in conds
        assert_raises(ValueError, plot_compare_evokeds,
                      evoked, ci='fake')  # ci must be float or None
        contrast["red/stim"] = red
        contrast["blue/stim"] = blue
        plot_compare_evokeds(contrast, picks=[0], colors=['r', 'b'], ymax=0)

        # Hack to test plotting of maxfiltered data
        evoked_sss = evoked.copy()
        evoked_sss.info['proc_history'] = [dict(max_info=None)]
        evoked_sss.plot_white(cov)
        evoked_sss.plot_white(cov_fname)
        plt.close('all')
    evoked.plot_sensors()  # Test plot_sensors
    plt.close('all')

run_tests_if_main()
<commit_msg>and more coverage ...<commit_after># Authors: Alexandre Gramfort <alexandre.gramfort@telecom-paristech.fr>
#          Denis Engemann <denis.engemann@gmail.com>
#          Martin Luessi <mluessi@nmr.mgh.harvard.edu>
#          Eric Larson <larson.eric.d@gmail.com>
#          Cathy Nangini <cnangini@gmail.com>
#          Mainak Jas <mainak@neuro.hut.fi>
#          Jona Sassenhagen <jona.sassenhagen@gmail.com>
#
# License: Simplified BSD

import os.path as op
import warnings

import numpy as np
from numpy.testing import assert_raises

from mne import io, read_events, Epochs, pick_types, read_cov
from mne.channels import read_layout
from mne.utils import slow_test, run_tests_if_main
from mne.viz.evoked import _butterfly_onselect, plot_compare_evokeds
from mne.viz.utils import _fake_click

# Set our plotters to test mode
import matplotlib
matplotlib.use('Agg')  # for testing don't use X server

warnings.simplefilter('always')  # enable b/c these tests throw warnings

base_dir = op.join(op.dirname(__file__), '..', '..', 'io', 'tests', 'data')
evoked_fname = op.join(base_dir, 'test-ave.fif')
raw_fname = op.join(base_dir, 'test_raw.fif')
cov_fname = op.join(base_dir, 'test-cov.fif')
event_name = op.join(base_dir, 'test-eve.fif')
event_id, tmin, tmax = 1, -0.1, 0.1
n_chan = 6
layout = read_layout('Vectorview-all')

def _get_raw():
    return io.read_raw_fif(raw_fname, preload=False)

def _get_events():
    return read_events(event_name)

def _get_picks(raw):
    return pick_types(raw.info, meg=True, eeg=False, stim=False,
                      ecg=False, eog=False, exclude='bads')

def _get_epochs():
    raw = _get_raw()
    raw.add_proj([], remove_existing=True)
    events = _get_events()
    picks = _get_picks(raw)
    # Use a subset of channels for plotting speed
    picks = picks[np.round(np.linspace(0, len(picks) - 1, n_chan)).astype(int)]
    picks[0] = 2  # make sure we have a magnetometer
    with warnings.catch_warnings(record=True):  # proj
        epochs = Epochs(raw, events[:5], event_id, tmin, tmax, picks=picks,
                        baseline=(None, 0))
    epochs.info['bads'] = [epochs.ch_names[-1]]
    return epochs

def _get_epochs_delayed_ssp():
    raw = _get_raw()
    events = _get_events()
    picks = _get_picks(raw)
    reject = dict(mag=4e-12)
    epochs_delayed_ssp = Epochs(raw, events[:10], event_id, tmin, tmax,
                                picks=picks, baseline=(None, 0),
                                proj='delayed', reject=reject)
    return epochs_delayed_ssp

@slow_test
def test_plot_evoked():
    """Test plotting of evoked
    """
    import matplotlib.pyplot as plt
    evoked = _get_epochs().average()
    with warnings.catch_warnings(record=True):
        fig = evoked.plot(proj=True, hline=[1], exclude=[], window_title='foo')
        # Test a click
        ax = fig.get_axes()[0]
        line = ax.lines[0]
        _fake_click(fig, ax,
                    [line.get_xdata()[0], line.get_ydata()[0]], 'data')
        _fake_click(fig, ax,
                    [ax.get_xlim()[0], ax.get_ylim()[1]], 'data')
        # plot with bad channels excluded & spatial_colors & zorder
        evoked.plot(exclude='bads')
        evoked.plot(exclude=evoked.info['bads'], spatial_colors=True, gfp=True,
                    zorder='std')

        # test selective updating of dict keys is working.
        evoked.plot(hline=[1], units=dict(mag='femto foo'))
        evoked_delayed_ssp = _get_epochs_delayed_ssp().average()
        evoked_delayed_ssp.plot(proj='interactive')
        evoked_delayed_ssp.apply_proj()
        assert_raises(RuntimeError, evoked_delayed_ssp.plot,
                      proj='interactive')
        evoked_delayed_ssp.info['projs'] = []
        assert_raises(RuntimeError, evoked_delayed_ssp.plot,
                      proj='interactive')
        assert_raises(RuntimeError, evoked_delayed_ssp.plot,
                      proj='interactive', axes='foo')
        plt.close('all')

        # test GFP only
        evoked.plot(gfp='only')
        assert_raises(ValueError, evoked.plot, gfp='foo')

        evoked.plot_image(proj=True)
        # plot with bad channels excluded
        evoked.plot_image(exclude='bads', cmap='interactive')
        evoked.plot_image(exclude=evoked.info['bads'])  # does the same thing
        plt.close('all')

        evoked.plot_topo()  # should auto-find layout
        _butterfly_onselect(0, 200, ['mag', 'grad'], evoked)
        plt.close('all')

        cov = read_cov(cov_fname)
        cov['method'] = 'empirical'
        evoked.plot_white(cov)
        evoked.plot_white([cov, cov])

        # plot_compare_evokeds: test condition contrast, CI, color assignment
        plot_compare_evokeds(evoked.copy().pick_types(meg='mag'))
        evoked.rename_channels({'MEG 2142': "MEG 1642"})
        colors = dict(red='r', blue='b')
        linestyles = dict(red='--', blue='-')
        red, blue = evoked.copy(), evoked.copy()
        red.data *= 1.1
        blue.data *= 0.9
        plot_compare_evokeds([red, blue], picks=3)  # list of evokeds
        plot_compare_evokeds([[red, evoked], [blue, evoked]])  # list of lists
        # test picking & plotting grads
        contrast = dict()
        contrast["red/stim"] = list((evoked.copy(), red))
        contrast["blue/stim"] = list((evoked.copy(), blue))
        # test a bunch of params at once
        plot_compare_evokeds(contrast, colors=colors, linestyles=linestyles,
                             picks=[0, 2], vlines=[.01, -.04], invert_y=True,
                             truncate_yaxis=False, ymax=10, ymin=-10,
                             styles={"red/stim": {"linewidth": 1}})
        assert_raises(ValueError, plot_compare_evokeds,
                      contrast, picks=[0, 3])  # bad picks: multiple types
        assert_raises(ValueError, plot_compare_evokeds,
                      contrast, picks='str')  # bad picks: not int or 'gfp'
        assert_raises(ValueError, plot_compare_evokeds,
                      evoked, colors=dict(fake=1))  # 'fake' not in conds
        assert_raises(ValueError, plot_compare_evokeds,
                      evoked, styles=dict(fake=1))  # 'fake' not in conds
        assert_raises(ValueError, plot_compare_evokeds,
                      evoked, ci='fake')  # ci must be float or None
        assert_raises(ValueError, plot_compare_evokeds, dict(1=1))
        contrast["red/stim"] = red
        contrast["blue/stim"] = blue
        plot_compare_evokeds(contrast, picks=[0], colors=['r', 'b'], ymax=0)

        # Hack to test plotting of maxfiltered data
        evoked_sss = evoked.copy()
        evoked_sss.info['proc_history'] = [dict(max_info=None)]
        evoked_sss.plot_white(cov)
        evoked_sss.plot_white(cov_fname)
        plt.close('all')
    evoked.plot_sensors()  # Test plot_sensors
    plt.close('all')

run_tests_if_main()
<|endoftext|>


<commit_before># Copyright (c) Metakernel Development Team.
# Distributed under the terms of the Modified BSD License.

from metakernel import Magic

class HelpMagic(Magic):

    def help_strings(self):
        suffixes = [
            "item{0}{0} - get detailed, technical information on item",
            "item{0}  - get help on item",
        ]
        prefixes = [
            "{0}{0}item - get detailed, technical information on item",
            "{0}item  - get help on item",
        ]
        strings = []
        if self.kernel.help_suffix:
            strings += [s.format(self.kernel.help_suffix['help'])
                        for s in suffixes]
        if 'help' in self.kernel.magic_prefixes:
            strings += [p.format(self.kernel.magic_prefixes['help'])
                        for p in prefixes]
        return sorted(strings)

    def line_help(self, text):
        """
        %help TEXT - get help on the given text

        This line magic shows the help for TEXT. Shows the help in the
        system pager that opens at bottom of screen.

        Example:
            %help dir

        """
        return self.get_help_on(text, 0, False)

    def cell_help(self, text):
        """
        %%help TEXT - get detailed help on the given text

        This line magic looks like a cell magic, but it really gets
        more detailed help on TEXT. Shows the help in the system
        pager that opens at bottom of screen.

        Example:
           %%help dir

        """
        return self.get_help_on(text, 1, False)

    def get_help_on(self, text, level, none_on_fail=False):
        """
        Examples
        ========

        All of the following give help on np.ones using the %python magic:

        %help %python np.ones
        %python np.ones?

        To get help on a magic itself, we would write one of the following:

        %python?
        %help %python
        ?%python
        """
        text = self._prep_text(text)

        info = self.kernel.parse_code(text)

        if info['magic']:

            minfo = info['magic']
            errmsg = "No such %s magic '%s'" % (minfo['type'], minfo['name'])

            if minfo['type'] == 'line':
                magic = self.kernel.line_magics.get(minfo['name'], None)

            else:
                magic = self.kernel.cell_magics.get(minfo['name'], None)

                if minfo['code']:
                    info = self.kernel.parse_code(minfo['code'])
                else:
                    info = self.kernel.parse_code(minfo['args'])

            if magic:
                info = self.kernel.parse_code(minfo['args'])
                return magic.get_help_on(info)

            elif not info['magic']['name']:
                return self.kernel.get_usage()

            else:
                return errmsg

        else:
            return self.kernel.get_kernel_help_on(info, level, none_on_fail)

    def _prep_text(self, text):
        text = text.strip()

        magic = self.kernel.magic_prefixes['magic']
        prefix = self.kernel.magic_prefixes['help']
        suffix = self.kernel.help_suffix

        text = text.replace('{0}{0}help'.format(magic), '')
        text = text.replace('{0}help'.format(magic), '')

        if text.startswith('{0}{0}'.format(prefix)):
            text = text[2:]
        elif text.startswith('{0}'.format(prefix)):
            text = text[1:]

        if suffix:
            if text.endswith('{0}{0}'.format(suffix)):
                text = text[:-2]
            elif text.endswith('{0}'.format(suffix)):
                text = text[:-1]

        return text

def register_magics(kernel):
    kernel.register_magics(HelpMagic)
<commit_msg>Fix handling of 'code' and 'args' in help.<commit_after># Copyright (c) Metakernel Development Team.
# Distributed under the terms of the Modified BSD License.

from metakernel import Magic

class HelpMagic(Magic):

    def help_strings(self):
        suffixes = [
            "item{0}{0} - get detailed, technical information on item",
            "item{0}  - get help on item",
        ]
        prefixes = [
            "{0}{0}item - get detailed, technical information on item",
            "{0}item  - get help on item",
        ]
        strings = []
        if self.kernel.help_suffix:
            strings += [s.format(self.kernel.help_suffix['help'])
                        for s in suffixes]
        if 'help' in self.kernel.magic_prefixes:
            strings += [p.format(self.kernel.magic_prefixes['help'])
                        for p in prefixes]
        return sorted(strings)

    def line_help(self, text):
        """
        %help TEXT - get help on the given text

        This line magic shows the help for TEXT. Shows the help in the
        system pager that opens at bottom of screen.

        Example:
            %help dir

        """
        return self.get_help_on(text, 0, False)

    def cell_help(self, text):
        """
        %%help TEXT - get detailed help on the given text

        This line magic looks like a cell magic, but it really gets
        more detailed help on TEXT. Shows the help in the system
        pager that opens at bottom of screen.

        Example:
           %%help dir

        """
        return self.get_help_on(text, 1, False)

    def get_help_on(self, text, level, none_on_fail=False):
        """
        Examples
        ========

        All of the following give help on np.ones using the %python magic:

        %help %python np.ones
        %python np.ones?

        To get help on a magic itself, we would write one of the following:

        %python?
        %help %python
        ?%python
        """
        text = self._prep_text(text)

        info = self.kernel.parse_code(text)

        if info['magic']:

            minfo = info['magic']
            errmsg = "No such %s magic '%s'" % (minfo['type'], minfo['name'])

            if minfo['type'] == 'line':
                magic = self.kernel.line_magics.get(minfo['name'], None)
                info = self.kernel.parse_code(minfo['args'])

            else:
                magic = self.kernel.cell_magics.get(minfo['name'], None)

                if minfo['code']:
                    info = self.kernel.parse_code(minfo['code'])
                else:
                    info = self.kernel.parse_code(minfo['args'])

            if magic:
                return magic.get_help_on(info)

            elif not info['magic']['name']:
                return self.kernel.get_usage()

            else:
                return errmsg

        else:
            return self.kernel.get_kernel_help_on(info, level, none_on_fail)

    def _prep_text(self, text):
        text = text.strip()

        magic = self.kernel.magic_prefixes['magic']
        prefix = self.kernel.magic_prefixes['help']
        suffix = self.kernel.help_suffix

        text = text.replace('{0}{0}help'.format(magic), '')
        text = text.replace('{0}help'.format(magic), '')

        if text.startswith('{0}{0}'.format(prefix)):
            text = text[2:]
        elif text.startswith('{0}'.format(prefix)):
            text = text[1:]

        if suffix:
            if text.endswith('{0}{0}'.format(suffix)):
                text = text[:-2]
            elif text.endswith('{0}'.format(suffix)):
                text = text[:-1]

        return text

def register_magics(kernel):
    kernel.register_magics(HelpMagic)
<|endoftext|>


<commit_before>
#  IATI Data Quality, tools for Data QA on IATI-formatted  publications
#  by Mark Brough, Martin Keegan, Ben Webb and Jennifer Smith
#
#  Copyright (C) 2013  Publish What You Fund
#
#  This programme is free software; you may redistribute and/or modify
#  it under the terms of the GNU Affero General Public License v3.0

from flask import Flask, render_template, flash, request, Markup, \
    session, redirect, url_for, escape, Response, abort, send_file
from flask.ext.login import (LoginManager, current_user, login_required,
                            login_user, logout_user, UserMixin, AnonymousUser,
                            confirm_login, fresh_login_required)

from iatidataquality import app
from iatidataquality import db

import os
import sys

current = os.path.dirname(os.path.abspath(__file__))
parent = os.path.dirname(current)
sys.path.append(parent)

from iatidq import dqindicators, dqsurveys, dqorganisations

from iatidq.models import *

import unicodecsv
import json
import json

@app.route("/surveys/admin/")
def surveys_admin():
    surveys = dqsurveys.surveys()
    workflows = dqsurveys.workflows()
    publishedstatuses=dqsurveys.publishedStatus()
    return render_template("surveys/surveys_admin.html",
                           workflows=workflows,
                           publishedstatuses=publishedstatuses,
                           surveys=surveys)

@app.route("/surveys/create/", methods=["GET", "POST"])
@app.route("/surveys/<organisation_code>/create/", methods=["GET", "POST"])
def create_survey(organisation_code=None):
    return "You're trying to create a survey"

def get_old_organisation_id(organisation_code='GB-1'):
    path = app.config["DATA_STORAGE_DIR"]
    old_organisation_file = os.path.join(path, '2012_2013_organisation_mapping.csv')

    old_organisation_data = unicodecsv.DictReader(file(old_organisation_file))
    for row in old_organisation_data:
        if row['2013_id'] == organisation_code:
            return row['2012_id']

def get_old_indicators():
    path = app.config["DATA_STORAGE_DIR"]
    old_indicators_file = os.path.join(path, '2012_indicators.csv')
    old_indicators_data = unicodecsv.DictReader(file(old_indicators_file))

    indicator_data = {}
    for row in old_indicators_data:
        if ((row["question_number"]) and (row["2013_indicator_name"])):
            indicator_data[int(row["question_number"])] = row["2013_indicator_name"]
    return indicator_data

def get_organisation_results(organisation_code, newindicators):
    old_organisation_id = get_old_organisation_id(organisation_code)
    indicators = get_old_indicators()

    path = app.config["DATA_STORAGE_DIR"]

    old_results_file = os.path.join(path, '2012_results.csv')
    old_results_data = unicodecsv.DictReader(file(old_results_file))

    data = {}

    for d in old_results_data:
        if (d["target_id"] == old_organisation_id):
            try:
                question_id = int(d["question_id"])
                d["newindicator_id"] = indicators[(int(d["question_id"]))]
                data[indicators[(int(d["question_id"]))]] = d
            except KeyError:
                pass
    for indicator in newindicators:
        try:
            print data[indicator.name]
        except KeyError:
            data[indicator.name] = {
                'result': ''
            }
    return data

@app.route("/organisations/<organisation_code>/survey/")
def organisation_survey(organisation_code=None):
    organisation = dqorganisations.organisations(organisation_code)
    survey = dqsurveys.getSurvey(organisation_code)
    workflows = dqsurveys.workflows()
    return render_template("surveys/survey.html",
                           organisation=organisation,
                           survey=survey,
                           workflows=workflows)

def getTimeRemainingNotice(deadline):
    time_remaining = ((deadline.date())-(datetime.utcnow().date()))
    if time_remaining.days >1:
        time_remaining_notice = "You have " + str(time_remaining.days) + " days to submit your response."
    else:
        time_remaining_notice = "Today is the last day for making any changes to your survey."
    return time_remaining_notice

def _survey_process_collect(organisation, workflow, request, organisationsurvey):
    indicators = request.form.getlist('indicator')

    workflow_id = workflow.Workflow.id
    currentworkflow_deadline = organisationsurvey.currentworkflow_deadline

    for indicator in indicators:
        data = {
            'organisationsurvey_id': organisationsurvey.id,
            'indicator_id': indicator,
            'workflow_id': workflow_id,
            'published_status': request.form.get(indicator+"-published"),
            'published_source': request.form.get(indicator+"-source"),
            'published_comment': request.form.get(indicator+"-comments"),
            'published_accepted': None
        }
        surveydata = dqsurveys.addSurveyData(data)

    if 'submit' in request.form:
        if workflow.Workflow.id == organisationsurvey.currentworkflow_id:
        # save data, change currentworkflow_id to leadsto
            dqsurveys.advanceSurvey(organisationsurvey)
            flash('Successfully submitted survey data', 'success')
        else:
            flash("Your survey data was updated.", 'warning')
    else:
        time_remaining_notice = getTimeRemainingNotice(organisationsurvey.currentworkflow_deadline)

        flash('Note: your survey has not yet been submitted. '+ time_remaining_notice, 'warning')

def _survey_process_review(organisation, workflow, request, organisationsurvey):
    indicators = request.form.getlist('indicator')
    workflow_id = workflow.Workflow.id
    currentworkflow_deadline = organisationsurvey.currentworkflow_deadline

    for indicator in indicators:
        data = {
            'organisationsurvey_id': organisationsurvey.id,
            'indicator_id': indicator,
            'workflow_id': workflow_id,
            'published_status': request.form.get(indicator+"-published"),
            'published_source': request.form.get(indicator+"-source"),
            'published_comment': request.form.get(indicator+"-comments"),
            'published_accepted': None
        }
        surveydata = dqsurveys.addSurveyData(data)

    if 'submit' in request.form:
        if workflow.Workflow.id == organisationsurvey.currentworkflow_id:
        # save data, change currentworkflow_id to leadsto
            dqsurveys.advanceSurvey(organisationsurvey)
            flash('Successfully submitted survey data', 'success')
        else:
            flash("Your survey data was updated.", 'warning')
    else:
        time_remaining_notice = getTimeRemainingNotice(organisationsurvey.currentworkflow_deadline)

        flash('Note: your survey has not yet been submitted. '+ time_remaining_notice, 'warning')

def _survey_process_finalreview(organisation, workflow, request, organisationsurvey):
    indicators = request.form.getlist('indicator')
    workflow_id = workflow.Workflow.id
    currentworkflow_deadline = organisationsurvey.currentworkflow_deadline

    for indicator in indicators:
        data = {
            'organisationsurvey_id': organisationsurvey.id,
            'indicator_id': indicator,
            'workflow_id': workflow_id,
            'published_status': request.form.get(indicator+"-published"),
            'published_source': request.form.get(indicator+"-source"),
            'published_comment': request.form.get(indicator+"-comments"),
            'published_accepted': request.form.get(indicator+"-agree")
        }
        surveydata = dqsurveys.addSurveyData(data)

    if 'submit' in request.form:
        if workflow.Workflow.id == organisationsurvey.currentworkflow_id:
        # save data, change currentworkflow_id to leadsto
            dqsurveys.advanceSurvey(organisationsurvey)
            flash('Successfully submitted survey data', 'success')
        else:
            flash("Your survey data was updated.", 'warning')
    else:
        time_remaining_notice = getTimeRemainingNotice(organisationsurvey.currentworkflow_deadline)

        flash('Note: your survey has not yet been submitted. '+ time_remaining_notice, 'warning')

def _survey_process_comment(organisation, workflow, request, organisationsurvey):
    indicators = request.form.getlist('indicator')
    workflow_id = workflow.Workflow.id
    currentworkflow_deadline = organisationsurvey.currentworkflow_deadline

    for indicator in indicators:
        data = {
            'organisationsurvey_id': organisationsurvey.id,
            'indicator_id': indicator,
            'workflow_id': workflow_id,
            'published_status': None,
            'published_source': None,
            'published_comment': request.form.get(indicator+"-comments"),
            'published_accepted': request.form.get(indicator+"-agree")
        }
        surveydata = dqsurveys.addSurveyData(data)

    if 'submit' in request.form:
        if workflow.Workflow.id == organisationsurvey.currentworkflow_id:
        # save data, change currentworkflow_id to leadsto
            dqsurveys.advanceSurvey(organisationsurvey)
            flash('Successfully submitted survey data', 'success')
        else:
            flash("Your survey data was updated.", 'warning')
    else:
        time_remaining_notice = getTimeRemainingNotice(organisationsurvey.currentworkflow_deadline)

        flash('Note: your survey has not yet been submitted. '+ time_remaining_notice, 'warning')

def _survey_process_send(organisation, workflow, request, organisationsurvey):
    indicators = request.form.getlist('indicator')

    #FIXME: need to actually send

    dqsurveys.advanceSurvey(organisationsurvey)
    flash('Successfully sent survey to donor.', 'success')

@app.route("/organisations/<organisation_code>/survey/<workflow_name>/", methods=["GET", "POST"])
def organisation_survey_edit(organisation_code=None, workflow_name=None):
    # FIXME: should probably go in setup
    dqsurveys.setupSurvey()

    workflow = dqsurveys.workflows(workflow_name)
    if not workflow:
        flash('That workflow does not exist.', 'error')
        return abort(404)

    organisation = dqorganisations.organisations(organisation_code)
    organisationsurvey = dqsurveys.getOrCreateSurvey({
                'organisation_id': organisation.id
                })

    if request.method=='POST':
        if (workflow.WorkflowType.name=='collect'):
            _survey_process_collect(organisation, workflow, request, organisationsurvey)
        elif workflow.WorkflowType.name=='send':
            if workflow.Workflow.id == organisationsurvey.currentworkflow_id:
                _survey_process_send(organisation_code, workflow, request, organisationsurvey)
            else:
                flash("Not possible to send survey to donor because it's not at the current stage in the workflow. Maybe you didn't submit the data, or maybe you already sent it to the donor?", 'error')
        elif workflow.WorkflowType.name=='review':
            _survey_process_review(organisation, workflow, request, organisationsurvey)
        elif workflow.WorkflowType.name=='comment':
            _survey_process_comment(organisation, workflow, request, organisationsurvey)
        elif workflow.WorkflowType.name=='finalreview':
            _survey_process_finalreview(organisation, workflow, request, organisationsurvey)
        elif workflow.WorkflowType.name=='finalised':
            return "finalised"
        return redirect(url_for("organisations", organisation_code=organisation_code))
    else:
        organisation = Organisation.query.filter_by(
            organisation_code=organisation_code).first_or_404()

        surveydata = dqsurveys.getSurveyData(organisation_code, workflow_name)
        surveydata_allworkflows = dqsurveys.getSurveyDataAllWorkflows(organisation_code)

        indicators = dqindicators.indicators("pwyf2013")
        org_indicators = dqorganisations._organisation_indicators_split(
        organisation, 2)["zero"]
        twentytwelvedata=get_organisation_results(organisation_code, indicators)
        publishedstatuses = dqsurveys.publishedStatus()
        publishedstatuses = dict(map(lambda ps: (ps.id, ps), publishedstatuses))

        old_publication_status= {
            '4': {
                'text': 'Always published',
                'class': 'success'
            },
            '3':  {
                'text': 'Sometimes published',
                'class': 'warning'
            },
            '2':  {
                'text': 'Collected',
                'class': 'important'
            },
            '1':  {
                'text': 'Not collected',
                'class': 'inverse'
            },
            '0':  {
                'text': 'Unknown',
                'class': ''
            },
            "": {
                'text': '',
                'class': ''
            }
        }
        template_path = "surveys/_survey_"+workflow.WorkflowType.name+".html"
        return render_template(template_path,
           organisation=organisation,
           indicators=indicators,
           org_indicators = org_indicators,
           twentytwelvedata=twentytwelvedata,
           old_publication_status=old_publication_status,
           publishedstatuses=publishedstatuses,
           workflow=workflow,
           surveydata=surveydata_allworkflows,
           organisationsurvey=organisationsurvey)
<commit_msg>Adjust dependencies<commit_after>
#  IATI Data Quality, tools for Data QA on IATI-formatted  publications
#  by Mark Brough, Martin Keegan, Ben Webb and Jennifer Smith
#
#  Copyright (C) 2013  Publish What You Fund
#
#  This programme is free software; you may redistribute and/or modify
#  it under the terms of the GNU Affero General Public License v3.0

from flask import Flask, render_template, flash, request, Markup, \
    session, redirect, url_for, escape, Response, abort, send_file

from iatidataquality import app
from iatidataquality import db

import os
import sys
import unicodecsv

current = os.path.dirname(os.path.abspath(__file__))
parent = os.path.dirname(current)
sys.path.append(parent)

from iatidq import dqindicators, dqsurveys, dqorganisations
from iatidq.models import *

import usermanagement

@app.route("/surveys/admin/")
def surveys_admin():
    surveys = dqsurveys.surveys()
    workflows = dqsurveys.workflows()
    publishedstatuses=dqsurveys.publishedStatus()
    return render_template("surveys/surveys_admin.html",
                           workflows=workflows,
                           publishedstatuses=publishedstatuses,
                           surveys=surveys)

@app.route("/surveys/create/", methods=["GET", "POST"])
@app.route("/surveys/<organisation_code>/create/", methods=["GET", "POST"])
def create_survey(organisation_code=None):
    return "You're trying to create a survey"

def get_old_organisation_id(organisation_code='GB-1'):
    path = app.config["DATA_STORAGE_DIR"]
    old_organisation_file = os.path.join(path, '2012_2013_organisation_mapping.csv')

    old_organisation_data = unicodecsv.DictReader(file(old_organisation_file))
    for row in old_organisation_data:
        if row['2013_id'] == organisation_code:
            return row['2012_id']

def get_old_indicators():
    path = app.config["DATA_STORAGE_DIR"]
    old_indicators_file = os.path.join(path, '2012_indicators.csv')
    old_indicators_data = unicodecsv.DictReader(file(old_indicators_file))

    indicator_data = {}
    for row in old_indicators_data:
        if ((row["question_number"]) and (row["2013_indicator_name"])):
            indicator_data[int(row["question_number"])] = row["2013_indicator_name"]
    return indicator_data

def get_organisation_results(organisation_code, newindicators):
    old_organisation_id = get_old_organisation_id(organisation_code)
    indicators = get_old_indicators()

    path = app.config["DATA_STORAGE_DIR"]

    old_results_file = os.path.join(path, '2012_results.csv')
    old_results_data = unicodecsv.DictReader(file(old_results_file))

    data = {}

    for d in old_results_data:
        if (d["target_id"] == old_organisation_id):
            try:
                question_id = int(d["question_id"])
                d["newindicator_id"] = indicators[(int(d["question_id"]))]
                data[indicators[(int(d["question_id"]))]] = d
            except KeyError:
                pass
    for indicator in newindicators:
        try:
            print data[indicator.name]
        except KeyError:
            data[indicator.name] = {
                'result': ''
            }
    return data

@app.route("/organisations/<organisation_code>/survey/")
def organisation_survey(organisation_code=None):
    organisation = dqorganisations.organisations(organisation_code)
    survey = dqsurveys.getSurvey(organisation_code)
    workflows = dqsurveys.workflows()
    return render_template("surveys/survey.html",
                           organisation=organisation,
                           survey=survey,
                           workflows=workflows)

def getTimeRemainingNotice(deadline):
    time_remaining = ((deadline.date())-(datetime.utcnow().date()))
    if time_remaining.days >1:
        time_remaining_notice = "You have " + str(time_remaining.days) + " days to submit your response."
    else:
        time_remaining_notice = "Today is the last day for making any changes to your survey."
    return time_remaining_notice

def _survey_process_collect(organisation, workflow, request, organisationsurvey):
    indicators = request.form.getlist('indicator')

    workflow_id = workflow.Workflow.id
    currentworkflow_deadline = organisationsurvey.currentworkflow_deadline

    for indicator in indicators:
        data = {
            'organisationsurvey_id': organisationsurvey.id,
            'indicator_id': indicator,
            'workflow_id': workflow_id,
            'published_status': request.form.get(indicator+"-published"),
            'published_source': request.form.get(indicator+"-source"),
            'published_comment': request.form.get(indicator+"-comments"),
            'published_accepted': None
        }
        surveydata = dqsurveys.addSurveyData(data)

    if 'submit' in request.form:
        if workflow.Workflow.id == organisationsurvey.currentworkflow_id:
        # save data, change currentworkflow_id to leadsto
            dqsurveys.advanceSurvey(organisationsurvey)
            flash('Successfully submitted survey data', 'success')
        else:
            flash("Your survey data was updated.", 'warning')
    else:
        time_remaining_notice = getTimeRemainingNotice(organisationsurvey.currentworkflow_deadline)

        flash('Note: your survey has not yet been submitted. '+ time_remaining_notice, 'warning')

def _survey_process_review(organisation, workflow, request, organisationsurvey):
    indicators = request.form.getlist('indicator')
    workflow_id = workflow.Workflow.id
    currentworkflow_deadline = organisationsurvey.currentworkflow_deadline

    for indicator in indicators:
        data = {
            'organisationsurvey_id': organisationsurvey.id,
            'indicator_id': indicator,
            'workflow_id': workflow_id,
            'published_status': request.form.get(indicator+"-published"),
            'published_source': request.form.get(indicator+"-source"),
            'published_comment': request.form.get(indicator+"-comments"),
            'published_accepted': None
        }
        surveydata = dqsurveys.addSurveyData(data)

    if 'submit' in request.form:
        if workflow.Workflow.id == organisationsurvey.currentworkflow_id:
        # save data, change currentworkflow_id to leadsto
            dqsurveys.advanceSurvey(organisationsurvey)
            flash('Successfully submitted survey data', 'success')
        else:
            flash("Your survey data was updated.", 'warning')
    else:
        time_remaining_notice = getTimeRemainingNotice(organisationsurvey.currentworkflow_deadline)

        flash('Note: your survey has not yet been submitted. '+ time_remaining_notice, 'warning')

def _survey_process_finalreview(organisation, workflow, request, organisationsurvey):
    indicators = request.form.getlist('indicator')
    workflow_id = workflow.Workflow.id
    currentworkflow_deadline = organisationsurvey.currentworkflow_deadline

    for indicator in indicators:
        data = {
            'organisationsurvey_id': organisationsurvey.id,
            'indicator_id': indicator,
            'workflow_id': workflow_id,
            'published_status': request.form.get(indicator+"-published"),
            'published_source': request.form.get(indicator+"-source"),
            'published_comment': request.form.get(indicator+"-comments"),
            'published_accepted': request.form.get(indicator+"-agree")
        }
        surveydata = dqsurveys.addSurveyData(data)

    if 'submit' in request.form:
        if workflow.Workflow.id == organisationsurvey.currentworkflow_id:
        # save data, change currentworkflow_id to leadsto
            dqsurveys.advanceSurvey(organisationsurvey)
            flash('Successfully submitted survey data', 'success')
        else:
            flash("Your survey data was updated.", 'warning')
    else:
        time_remaining_notice = getTimeRemainingNotice(organisationsurvey.currentworkflow_deadline)

        flash('Note: your survey has not yet been submitted. '+ time_remaining_notice, 'warning')

def _survey_process_comment(organisation, workflow, request, organisationsurvey):
    indicators = request.form.getlist('indicator')
    workflow_id = workflow.Workflow.id
    currentworkflow_deadline = organisationsurvey.currentworkflow_deadline

    for indicator in indicators:
        data = {
            'organisationsurvey_id': organisationsurvey.id,
            'indicator_id': indicator,
            'workflow_id': workflow_id,
            'published_status': None,
            'published_source': None,
            'published_comment': request.form.get(indicator+"-comments"),
            'published_accepted': request.form.get(indicator+"-agree")
        }
        surveydata = dqsurveys.addSurveyData(data)

    if 'submit' in request.form:
        if workflow.Workflow.id == organisationsurvey.currentworkflow_id:
        # save data, change currentworkflow_id to leadsto
            dqsurveys.advanceSurvey(organisationsurvey)
            flash('Successfully submitted survey data', 'success')
        else:
            flash("Your survey data was updated.", 'warning')
    else:
        time_remaining_notice = getTimeRemainingNotice(organisationsurvey.currentworkflow_deadline)

        flash('Note: your survey has not yet been submitted. '+ time_remaining_notice, 'warning')

def _survey_process_send(organisation, workflow, request, organisationsurvey):
    indicators = request.form.getlist('indicator')

    #FIXME: need to actually send

    dqsurveys.advanceSurvey(organisationsurvey)
    flash('Successfully sent survey to donor.', 'success')

@app.route("/organisations/<organisation_code>/survey/<workflow_name>/", methods=["GET", "POST"])
def organisation_survey_edit(organisation_code=None, workflow_name=None):
    # FIXME: should probably go in setup
    dqsurveys.setupSurvey()

    workflow = dqsurveys.workflows(workflow_name)
    if not workflow:
        flash('That workflow does not exist.', 'error')
        return abort(404)

    organisation = dqorganisations.organisations(organisation_code)
    organisationsurvey = dqsurveys.getOrCreateSurvey({
                'organisation_id': organisation.id
                })

    if request.method=='POST':
        if (workflow.WorkflowType.name=='collect'):
            _survey_process_collect(organisation, workflow, request, organisationsurvey)
        elif workflow.WorkflowType.name=='send':
            if workflow.Workflow.id == organisationsurvey.currentworkflow_id:
                _survey_process_send(organisation_code, workflow, request, organisationsurvey)
            else:
                flash("Not possible to send survey to donor because it's not at the current stage in the workflow. Maybe you didn't submit the data, or maybe you already sent it to the donor?", 'error')
        elif workflow.WorkflowType.name=='review':
            _survey_process_review(organisation, workflow, request, organisationsurvey)
        elif workflow.WorkflowType.name=='comment':
            _survey_process_comment(organisation, workflow, request, organisationsurvey)
        elif workflow.WorkflowType.name=='finalreview':
            _survey_process_finalreview(organisation, workflow, request, organisationsurvey)
        elif workflow.WorkflowType.name=='finalised':
            return "finalised"
        return redirect(url_for("organisations", organisation_code=organisation_code))
    else:
        organisation = Organisation.query.filter_by(
            organisation_code=organisation_code).first_or_404()

        surveydata = dqsurveys.getSurveyData(organisation_code, workflow_name)
        surveydata_allworkflows = dqsurveys.getSurveyDataAllWorkflows(organisation_code)

        indicators = dqindicators.indicators("pwyf2013")
        org_indicators = dqorganisations._organisation_indicators_split(
        organisation, 2)["zero"]
        twentytwelvedata=get_organisation_results(organisation_code, indicators)
        publishedstatuses = dqsurveys.publishedStatus()
        publishedstatuses = dict(map(lambda ps: (ps.id, ps), publishedstatuses))

        old_publication_status= {
            '4': {
                'text': 'Always published',
                'class': 'success'
            },
            '3':  {
                'text': 'Sometimes published',
                'class': 'warning'
            },
            '2':  {
                'text': 'Collected',
                'class': 'important'
            },
            '1':  {
                'text': 'Not collected',
                'class': 'inverse'
            },
            '0':  {
                'text': 'Unknown',
                'class': ''
            },
            "": {
                'text': '',
                'class': ''
            }
        }
        template_path = "surveys/_survey_"+workflow.WorkflowType.name+".html"
        return render_template(template_path,
           organisation=organisation,
           indicators=indicators,
           org_indicators = org_indicators,
           twentytwelvedata=twentytwelvedata,
           old_publication_status=old_publication_status,
           publishedstatuses=publishedstatuses,
           workflow=workflow,
           surveydata=surveydata_allworkflows,
           organisationsurvey=organisationsurvey)
<|endoftext|>


<commit_before>#!/usr/bin/env python
import re, glob, sys,os
import argparse
import itertools as it
import numpy as np
import scipy.sparse
import cPickle as pickle
from sklearn.naive_bayes import MultinomialNB
from content_sources import arxiv, bibtex, rss

def find_keywords(text):

    keywords=re.sub('[{}:?!@#$%^&*\(\)_.\\/,\'\"]','',text).upper()
    prepositions = open('data/prepositions.dat').read().upper().split()

    for p in prepositions:
        keywords = re.sub(r"\b{!s}\b".format(p),' ',keywords)

    return keywords.split()

def test(argv):
    #train the algorithm
    #read in titles we like
    good_titles_kw = find_keywords(open('data/allTitles.dat').read())
    #read in titles we don't
    bad_titles_kw = find_keywords(open('data/arxivTitles.500.dat').read())

    #remore duplicate words
    good_kw_set = set(good_titles_kw)
    bad_kw_set = set(bad_titles_kw)

    #make a set of all the words
    all_kw_set = good_kw_set.union(bad_kw_set)

    #count the occurences of each of the the keywords
    good_freq = np.zeros(len(all_kw_set))
    bad_freq = np.zeros(len(all_kw_set))
    all_freq = np.zeros(len(all_kw_set))

    ##good_txt = ' '.join(good_titles_kw)
    ##bad_txt = ' '.join(bad_titles_kw)

    for i,kw in enumerate(all_kw_set):
        good_freq[i] += good_titles_kw.count(kw)# len(re.findall(r"\b{!s}\b".format(kw), good_txt))
        bad_freq[i] += bad_titles_kw.count(kw) # len(re.findall(r"\b{!s}\b".format(kw), bad_txt))

    all_freq = good_freq+bad_freq

    #compute the base rates of words in the good or bad set
    base_good = np.sum(good_freq)/np.sum(all_freq)
    base_bad = np.sum(bad_freq)/np.sum(all_freq)

    #compute the "evidence probabilites"
    evidence_prob = all_freq/np.sum(all_freq)

    #compute likelihood of word being good and bad
    good_likelihood  = good_freq/all_freq
    bad_likelihood  = bad_freq/all_freq

    #make a dictionary to easily access likelihoods of our keywords
    gl_dict = {kw:lh for kw,lh in zip(all_kw_set,good_likelihood)}
    bl_dict = {kw:lh for kw,lh in zip(all_kw_set,bad_likelihood)}

    #read in new titles
    new_titles = open('data/arxivTitles.feb20.300.dat').readlines()
    nt_good_scores = np.empty(len(new_titles))
    nt_bad_scores = np.empty(len(new_titles))

    #compute score for each new title
    for i,nt in enumerate(new_titles):
        nt_kw_set = set(find_keywords(nt))
        nt_good_scores[i] = base_good*np.prod([gl_dict[kw] for kw in nt_kw_set if kw in gl_dict])
        nt_bad_scores[i] = base_bad*np.prod([bl_dict[kw] for kw in nt_kw_set if kw in bl_dict])

    for b,g,t in zip(nt_bad_scores,nt_good_scores,new_titles):
        print(b,g,'GOOD!' if g>b and g>0 else 'BAD!',t)

#read existing knowledge from a database
def parse_knowledge(knowledge_db):
    pass

# \param nb = MultinomialNB classifier
def filter_content(content,method,nb):

    #new content
    new_entries = [find_keywords(entry[method])  for entry in content]

    #compute vector for each new entry
    X = scipy.sparse.lil_matrix((len(good_samples) + len(bad_samples),len(all_kw)))

    for j,kw in enumerate(all_kw):
        for i,gs in enumerate(good_samples):
            if kw in gs:
                X[i,j]+=1

    categories = np.predict(X)
    return [e for c,e in zip(categories,content) if c =='good']

def get_content(sources):
    all_content = list()
    for src in sources:
        try:
            src.fetch()
        except:
            print("Fetch of content from {!r} has failed".format(src))

        content = None
        try:
            content = src.parse()
        except:
            print("parsing of content from {!r} has failed".format(src))

        if content:
            all_content += content

    return all_content

#training suite
def train(good_sources, bad_sources,method,naive_bayes=None,keywords=list()):
    #train the algorithm
    good_samples = [find_keywords(entry[method]) for entry in good_sources]
    bad_samples = [find_keywords(entry[method]) for entry in bad_sources]

    new_kws = set([kw for kw in it.chain(*good_samples)]+ [kw for kw in it.chain(*bad_samples)])

    all_kw = keywords + [nkw for nkw in new_kws if nkw not in keywords]

    X = scipy.sparse.lil_matrix((len(good_samples) + len(bad_samples),len(all_kw)))

    offset = len(good_samples)
    for j,kw in enumerate(all_kw):
        for i,gs in enumerate(good_samples):
            X[i,j]=gs.count(kw)
        for i,bs in enumerate(bad_samples):
            X[i+offset,j]=bs.count(kw)
    y = ['good']*len(good_samples) + ['bad']*len(bad_samples)

    if naive_bayes is None:
        nb = MultinomialNB()
        nb.fit(X,y)
    else:
        nb = naive_bayes
        nb.partial_fit(X,y,classes=['good','bad'])

    return nb, all_kw

def to_markdown(content):
    pass

def main(argv):

#add command line options for sources, output prefs, database of "good" keywords
    parser = argparse.ArgumentParser()
    parser.add_argument('-b','--bibtex', help='bibtex files to fetch',dest='bibfiles', nargs='*',default=list())
    parser.add_argument('-j','--journals', help='journals to fetch. Currently supports {}.'.format(' '.join(rss.rss_feeds.keys())),
                        nargs='*',dest='journals',default=list())
    parser.add_argument('-a','--arXiv', help='arXiv categories to fetch',
                      nargs='*',dest='arXiv',default=list())
    parser.add_argument('--all_sources', help='flag to search from all sources.',action ='store_true')
    parser.add_argument('--all_good_sources', help='flag to search from good sources. Specfied in your config file.',action ='store_true')
    parser.add_argument('--train', help='flag to train. All sources beside "--train-input-good" are treated as bad/irrelevant papers',action ='store_true')
    parser.add_argument('-g','--train_input_good', help='bibtex file containing relevant articles.',dest ='good_source',default=None)
    parser.add_argument('-m','--method', help='Methods to try to find relevent papers. Right now, only all, title, author, and abstract are valid fields',
                        dest='method',default='title')
    parser.add_argument('-k', '--knowledge',
                        help='path to database containing information about good and bad keywords. \
                              If you are training, you must specifiy this, as it will be where your output is written ',
                        dest='knowledge',default=None)
    parser.add_argument('--overwrite-knowledge', help='flag to overwrite knowledge,if training',action ='store_false',default=True, dest='overwrite_knowledge')

    args = parser.parse_args(argv)

    if  not args.method in ['title','abstract','author, all']:
            print("Invalid method. Options are title, abstract, author, and all")
            exit()

    method = args.method

    #Set up training if that's what we're doing
    if args.train:

        #check to make sure we have a good training input
        if args.good_source is None:
            print("When training, you must specify one good source")
            exit()
        if not os.path.exists(args.good_source):
            print("Specified training input does not exist")
            exit()
        if not os.path.isfile(args.good_source):
            print("Specified training input is not a file")
            exit()
        if not os.path.splitext(args.good_source)[1] == '.bib' :
            print("Training input must be in bibtex format")
            exit()

        #make sure you have a good destination to write your training output

        #existing naive_bayes object and keyword list
        nb=None
        kw=list()
        if args.knowledge is not None:
            if not os.path.isfile(args.knowledge):
                print("Specified training output exists and is not a file")
                exit()

            if not os.path.splitext(args.knowledge)[1] == '.p' :
                print("Training output must be a pickle file")
                exit()

            knowledge=args.knowledge
        else:
            knowledge =os.path.expanduser('~/.shakespeare')

        if os.path.exists(knowledge):
            if args.overwrite_knowledge:
                print("Overwriting existings knowledege database: {}".format(knowledge))
            else:
                print("Appending to existing knowledege database: {}".format(knowledge))
                nb=pickle.load(open(knowledge+'/nb.p'))
                kw=pickle.load(open(knowledge+'/kw.p'))

        good_content = get_content([bibtex.BibTex(args.good_source)])

        bad_content = get_content([arxiv.ArXiv(cat) for cat in args.arXiv] +
                                  [bibtex.BibTex(bibfile) for bibfile in args.bibfiles] +
                                  [rss.JournalFeed(journal) for journal in args.journals])

        #train, and write out knowledge (naive_bayes class and keywords)
        nb, kw = train(good_content,bad_content,method,naive_bayes=nb, keywords=kw)
        pickle.dump(nb,open(knowledge+'/nb.p','w'))
        pickle.dump(kw,open(knowledge+'/kw.p','w'))

    #we are filtering new content through our existing knowledge
    else:

        #locate our knowledge
        if (args.knowledge is None):
            pass

        nb = pickle.load(args.knowledge)

        sources = [arxiv.ArXiv(cat) for cat in args.arXiv] + \
                  [bibtex.BibTex (bibfile) for bibfile in args.bibfiles] + \
                  [rss.JournalFeed(journal) for j in journal]

        new_content = get_content(sources)

        relevant_content = filter_content(content,nb,method)

        to_markdown(relevant_content)

if __name__=="__main__":
    main(sys.argv[1:])
<commit_msg>finished methods to filter new content, write out markdown<commit_after>#!/usr/bin/env python
import re, glob, sys,os
import argparse
import itertools as it
import numpy as np
import scipy.sparse
import cPickle as pickle
from sklearn.naive_bayes import MultinomialNB
from content_sources import arxiv, bibtex, rss

def find_keywords(text):

    keywords=re.sub('[{}:?!@#$%^&*\(\)_.\\/,\'\"]','',text).upper()
    prepositions = open('data/prepositions.dat').read().upper().split()

    for p in prepositions:
        keywords = re.sub(r"\b{!s}\b".format(p),' ',keywords)

    return keywords.split()

def test(argv):
    #train the algorithm
    #read in titles we like
    good_titles_kw = find_keywords(open('data/allTitles.dat').read())
    #read in titles we don't
    bad_titles_kw = find_keywords(open('data/arxivTitles.500.dat').read())

    #remore duplicate words
    good_kw_set = set(good_titles_kw)
    bad_kw_set = set(bad_titles_kw)

    #make a set of all the words
    all_kw_set = good_kw_set.union(bad_kw_set)

    #count the occurences of each of the the keywords
    good_freq = np.zeros(len(all_kw_set))
    bad_freq = np.zeros(len(all_kw_set))
    all_freq = np.zeros(len(all_kw_set))

    ##good_txt = ' '.join(good_titles_kw)
    ##bad_txt = ' '.join(bad_titles_kw)

    for i,kw in enumerate(all_kw_set):
        good_freq[i] += good_titles_kw.count(kw)# len(re.findall(r"\b{!s}\b".format(kw), good_txt))
        bad_freq[i] += bad_titles_kw.count(kw) # len(re.findall(r"\b{!s}\b".format(kw), bad_txt))

    all_freq = good_freq+bad_freq

    #compute the base rates of words in the good or bad set
    base_good = np.sum(good_freq)/np.sum(all_freq)
    base_bad = np.sum(bad_freq)/np.sum(all_freq)

    #compute the "evidence probabilites"
    evidence_prob = all_freq/np.sum(all_freq)

    #compute likelihood of word being good and bad
    good_likelihood  = good_freq/all_freq
    bad_likelihood  = bad_freq/all_freq

    #make a dictionary to easily access likelihoods of our keywords
    gl_dict = {kw:lh for kw,lh in zip(all_kw_set,good_likelihood)}
    bl_dict = {kw:lh for kw,lh in zip(all_kw_set,bad_likelihood)}

    #read in new titles
    new_titles = open('data/arxivTitles.feb20.300.dat').readlines()
    nt_good_scores = np.empty(len(new_titles))
    nt_bad_scores = np.empty(len(new_titles))

    #compute score for each new title
    for i,nt in enumerate(new_titles):
        nt_kw_set = set(find_keywords(nt))
        nt_good_scores[i] = base_good*np.prod([gl_dict[kw] for kw in nt_kw_set if kw in gl_dict])
        nt_bad_scores[i] = base_bad*np.prod([bl_dict[kw] for kw in nt_kw_set if kw in bl_dict])

    for b,g,t in zip(nt_bad_scores,nt_good_scores,new_titles):
        print(b,g,'GOOD!' if g>b and g>0 else 'BAD!',t)

#read existing knowledge from a database
def parse_knowledge(knowledge_db):
    pass

# \param nb = MultinomialNB classifier
def filter_content(content,method,naive_bayes,keywords):

    new_samples = [find_keywords(entry[method])  for entry in content]

    #compute vector for each new entry
    print(len(new_samples),len(keywords))
    X = scipy.sparse.lil_matrix((len(new_samples),len(keywords)))

    for j,kw in enumerate(keywords):
        for i,gs in enumerate(new_samples):
            if kw in gs:
                X[i,j]+=1

    categories = naive_bayes.predict(X)
    return [e for c,e in zip(categories,content) if c =='good']

def get_content(sources):
    all_content = list()
    for src in sources:
        try:
            src.fetch()
        except:
            print("Fetch of content from {!r} has failed".format(src))

        content = None
        try:
            content = src.parse()
        except:
            print("parsing of content from {!r} has failed".format(src))

        if content:
            all_content += content

    return all_content

def load_knowledge(knowledge):
    #existing naive_bayes object and keyword list
    nb=None
    kw=list()
    if knowledge is not None:
        if not os.path.isdir(knowledge):
            print("Knowledge bust be a directory")
            exit()

        kfiles = glob.glob(knowledge+'/*')
        if not 'nb.p' in kfiles:
            print("Knowledge does not contain nb.p (pickled naive bayes object)")
            exit()
        if not 'kw.p' in kfiles:
            print("Knowledge does not contain kw.p (pickled keyword list)")
            exit()

    else:
        knowledge =os.path.expanduser('~/.shakespeare')

    if os.path.exists(knowledge):
        nb=pickle.load(open(knowledge+'/nb.p'))
        kw=pickle.load(open(knowledge+'/kw.p'))

    return(nb,kw)

#training suite
def train(good_sources, bad_sources,method,naive_bayes=None,keywords=list()):
    #train the algorithm
    good_samples = [find_keywords(entry[method]) for entry in good_sources]
    bad_samples = [find_keywords(entry[method]) for entry in bad_sources]

    new_kws = set([kw for kw in it.chain(*good_samples)]+ [kw for kw in it.chain(*bad_samples)])

    all_kw = keywords + [nkw for nkw in new_kws if nkw not in keywords]

    X = scipy.sparse.lil_matrix((len(good_samples) + len(bad_samples),len(all_kw)))

    offset = len(good_samples)
    for j,kw in enumerate(all_kw):
        for i,gs in enumerate(good_samples):
            X[i,j]=gs.count(kw)
        for i,bs in enumerate(bad_samples):
            X[i+offset,j]=bs.count(kw)
    y = ['good']*len(good_samples) + ['bad']*len(bad_samples)

    if naive_bayes is None:
        nb = MultinomialNB()
        nb.fit(X,y)
    else:
        nb = naive_bayes
        nb.partial_fit(X,y,classes=['good','bad'])

    return nb, all_kw

def to_markdown(content,output_file):
    try:
        with open(output_file,'w') as outf:
            outf.write('# Relevant articles\n')
            for article in content:
                outf.write("## {}\n".format(article['title']))
                outf.write("* authors: {}\n".format(article['author']))
                outf.write("* abstract: {}\n".format(article['abstract']))
                outf.write("* [link]({})\n\n".format(article['url']))
    except:
        print("Failed to write markdown file")

def main(argv):

    #add command line options for sources, output prefs, database of "good" keywords
    parser = argparse.ArgumentParser()
    parser.add_argument('-o','--output', help='output file name. only supports markdown right now.',dest ='output',default=None)
    parser.add_argument('-b','--bibtex', help='bibtex files to fetch',dest='bibfiles', nargs='*',default=list())
    parser.add_argument('-j','--journals', help='journals to fetch. Currently supports {}.'.format(' '.join(rss.rss_feeds.keys())),
                        nargs='*',dest='journals',default=list())
    parser.add_argument('-a','--arXiv', help='arXiv categories to fetch',
                      nargs='*',dest='arXiv',default=list())
    parser.add_argument('--all_sources', help='flag to search from all sources.',action ='store_true')
    parser.add_argument('--all_good_sources', help='flag to search from good sources. Specfied in your config file.',action ='store_true')
    parser.add_argument('--train', help='flag to train. All sources beside "--train-input-good" are treated as bad/irrelevant papers',action ='store_true')
    parser.add_argument('-g','--train_input_good', help='bibtex file containing relevant articles.',dest ='good_source',default=None)
    parser.add_argument('-m','--method', help='Methods to try to find relevent papers. Right now, only all, title, author, and abstract are valid fields',
                        dest='method',default='title')
    parser.add_argument('-k', '--knowledge',
                        help='path to database containing information about good and bad keywords. \
                              If you are training, you must specifiy this, as it will be where your output is written ',
                        dest='knowledge',default=None)
    parser.add_argument('--overwrite-knowledge', help='flag to overwrite knowledge,if training',action ='store_false',default=True, dest='overwrite_knowledge')

    parser.add_argument('--feedback', help='flag to overwrite knowledge,if training',action ='store_false',default=True, dest='overwrite_knowledge')
    args = parser.parse_args(argv)

    if  not args.method in ['title','abstract','author, all']:
            print("Invalid method. Options are title, abstract, author, and all")
            exit()

    method = args.method

    #Set up training if that's what we're doing
    if args.train:

        #check to make sure we have a good training input
        if args.good_source is None:
            print("When training, you must specify one good source")
            exit()
        if not os.path.exists(args.good_source):
            print("Specified training input does not exist")
            exit()
        if not os.path.isfile(args.good_source):
            print("Specified training input is not a file")
            exit()
        if not os.path.splitext(args.good_source)[1] == '.bib' :
            print("Training input must be in bibtex format")
            exit()

        #load the existing knowledge
        nb,kw = load_knowledge(args.knowledge)

        #make sure you have a good destination to write your training output

        good_content = get_content([bibtex.BibTex(args.good_source)])

        bad_content = get_content([arxiv.ArXiv(cat) for cat in args.arXiv] +
                                  [bibtex.BibTex(bibfile) for bibfile in args.bibfiles] +
                                  [rss.JournalFeed(journal) for journal in args.journals])

        #train, and write out knowledge (naive_bayes class and keywords)
        nb, kw = train(good_content,bad_content,method,naive_bayes=nb, keywords=kw)
        pickle.dump(nb,open(knowledge+'/nb.p','w'))
        pickle.dump(kw,open(knowledge+'/kw.p','w'))

    #we are filtering new content through our existing knowledge
    else:

        #load the old knowledge
        nb,kw = load_knowledge(args.knowledge)
        sources = [arxiv.ArXiv(cat) for cat in args.arXiv] + \
                  [bibtex.BibTex (bibfile) for bibfile in args.bibfiles] + \
                  [rss.JournalFeed(journal) for journal in args.journals]

        new_content = get_content(sources)

        relevant_content = filter_content(new_content,method,nb,kw)

        if (args.output):
            to_markdown(relevant_content,args.output)
        else:
            print(relevant_content)

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


<commit_before># Tai Sakuma <tai.sakuma@cern.ch>
from ..ProgressBar import NullProgressMonitor
import multiprocessing
from operator import itemgetter

##__________________________________________________________________||
class Worker(multiprocessing.Process):
    def __init__(self, task_queue, result_queue, lock, progressReporter):
        multiprocessing.Process.__init__(self)
        self.task_queue = task_queue
        self.result_queue = result_queue
        self.lock = lock
        self.progressReporter = progressReporter

    def run(self):
        while True:
            message = self.task_queue.get()
            if message is None:
                self.task_queue.task_done()
                break
            taskNo, task, args, kwargs = message
            try:
                result = task(progressReporter = self.progressReporter, *args, **kwargs)
            except TypeError:
                result = task(*args, **kwargs)
            self.task_queue.task_done()
            self.result_queue.put((taskNo, result))

##__________________________________________________________________||
class CommunicationChannel(object):
    """A communication channel with workers in other processes.

    You can send tasks to workers through this channel. The workers,
    running in other processes, execute the tasks in the background.
    You can receive the results of the tasks also through this
    channel.

    An instance of this class can be created with two optional
    arguments: ``nprocesses``, the number of workers to be created,
    and ``progressMonitor``::

        progressBar = ProgressBar()
        progressMonitor = BProgressMonitor(progressBar)
        channel = CommunicationChannel(nprocesses = 10, progressMonitor = progressMonitor)

    Workers will be created when ``begin()`` is called::

        channel.begin()

    In ``begin()``, this class gives each worker a
    ``progressReporter``, which is created by the ``progressMonitor``.

    Now, you are ready to send a task. A task is a function or any
    object which is callable and picklable and which takes the only
    argument ``progressReporter``. A value that a task returns is the
    result of the task and must be picklable. For example, an instance
    of ``EventLoop`` can be a task. You can send a task with the
    method ``put``::

        channel.put(task1)

    This class sends the task to a worker. The worker which receives
    the task will call the task with the ``progressReporter``.

    You can send multiple tasks::

        channel.put(task2)
        channel.put(task3)
        channel.put(task4)
        channel.put(task5)

    They will be executed by workers.

    You can receive the results of the tasks with the method
    ``receive()``::

        results = channel.receive()

    This method will wait until all tasks are finished. If a task
    gives a ``progressReport`` to the ``progressReporter``, the report
    will be used, for example, to update progress bars on the screen.

    When all tasks end, results will be returned. The return value
    ``results`` is a list of results of the tasks. They are sorted in
    the order in which the tasks are originally put.

    After receiving the results, you can put more tasks::

        channel.put(task6)
        channel.put(task7)

    And you can receive the results of them::

        results = channel.receive()

    If there are no more tasks to be done, you should call the method
    ``end``::

        channel.end()

    This will end all background processes.

    """

    def __init__(self, nprocesses = 16, progressMonitor = None):

        if nprocesses <= 0:
            raise ValueError("nprocesses must be at least one: " + str(nprocesses) + " is given")

        self.progressMonitor = NullProgressMonitor() if progressMonitor is None else progressMonitor
        self.nMaxProcesses = nprocesses
        self.nCurrentProcesses = 0
        self.task_queue = multiprocessing.JoinableQueue()
        self.result_queue = multiprocessing.Queue()
        self.lock = multiprocessing.Lock()
        self.nRunningTasks = 0
        self.taskNo = 0

    def begin(self):
        if self.nCurrentProcesses >= self.nMaxProcesses: return
        for i in xrange(self.nCurrentProcesses, self.nMaxProcesses):
            worker = Worker(
                task_queue = self.task_queue,
                result_queue = self.result_queue,
                progressReporter = self.progressMonitor.createReporter(),
                lock = self.lock
            )
            worker.start()
            self.nCurrentProcesses += 1

    def put(self, task, *args, **kwargs):
        self.task_queue.put((self.taskNo, task, args, kwargs))
        self.taskNo += 1
        self.nRunningTasks += 1

    def receive(self):
        messages = [ ] # a list of (taskNo, result)
        while self.nRunningTasks >= 1:
            if self.result_queue.empty(): continue
            message = self.result_queue.get()
            messages.append(message)
            self.nRunningTasks -= 1

        # sort in the order of taskNo
        messages = sorted(messages, key = itemgetter(0))

        results = [result for taskNo, result in messages]
        return results

    def end(self):
        for i in xrange(self.nCurrentProcesses):
            self.task_queue.put(None) # end workers
        self.task_queue.join()
        self.nCurrentProcesses = 0

##__________________________________________________________________||
<commit_msg>update docstrings in CommunicationChannel<commit_after># Tai Sakuma <tai.sakuma@cern.ch>
from ..ProgressBar import NullProgressMonitor
import multiprocessing
from operator import itemgetter

##__________________________________________________________________||
class Worker(multiprocessing.Process):
    def __init__(self, task_queue, result_queue, lock, progressReporter):
        multiprocessing.Process.__init__(self)
        self.task_queue = task_queue
        self.result_queue = result_queue
        self.lock = lock
        self.progressReporter = progressReporter

    def run(self):
        while True:
            message = self.task_queue.get()
            if message is None:
                self.task_queue.task_done()
                break
            taskNo, task, args, kwargs = message
            try:
                result = task(progressReporter = self.progressReporter, *args, **kwargs)
            except TypeError:
                result = task(*args, **kwargs)
            self.task_queue.task_done()
            self.result_queue.put((taskNo, result))

##__________________________________________________________________||
class CommunicationChannel(object):
    """A communication channel with workers in other processes.

    You can send tasks to workers through this channel. The workers,
    running in other processes, execute the tasks in the background.
    You can receive the results of the tasks also through this
    channel.

    An instance of this class can be created with two optional
    arguments: ``nprocesses``, the number of workers to be created,
    and ``progressMonitor``::

        progressBar = ProgressBar()
        progressMonitor = BProgressMonitor(progressBar)
        channel = CommunicationChannel(nprocesses = 10, progressMonitor = progressMonitor)

    Workers will be created when ``begin()`` is called::

        channel.begin()

    In ``begin()``, this class gives each worker a
    ``progressReporter``, which is created by the ``progressMonitor``.

    Now, you are ready to send a task. A task is a function or any
    object which is callable and picklable. A task can take any number
    of arguments. If a task takes a named argument
    ``progressReporter``, the worker will give the
    ``progressReporter`` to the task. A value that a task returns is
    the result of the task and must be picklable. For example, an
    instance of ``EventLoop`` can be a task. You can send a task with
    the method ``put``::

        channel.put(task1, 10, 20, A = 30)

    Here, 10, 20, A = 30 are the arguments to the task.

    This class sends the task to a worker. The worker which receives
    the task will first try to call the task with the
    ``progressReporter`` in addition to the arguments. If the task
    doesn't take the ``progressReporter``, it calls only with the
    arguments.

    You can send multiple tasks::

        channel.put(task2)
        channel.put(task3, 100, 200)
        channel.put(task4, A = 'abc')
        channel.put(task5)

    They will be executed by workers.

    You can receive the results of the tasks with the method
    ``receive()``::

        results = channel.receive()

    This method will wait until all tasks are finished. If a task
    gives a ``progressReport`` to the ``progressReporter``, the report
    will be used, for example, to update progress bars on the screen.

    When all tasks end, results will be returned. The return value
    ``results`` is a list of results of the tasks. They are sorted in
    the order in which the tasks are originally put.

    After receiving the results, you can put more tasks::

        channel.put(task6)
        channel.put(task7)

    And you can receive the results of them::

        results = channel.receive()

    If there are no more tasks to be done, you should call the method
    ``end``::

        channel.end()

    This will end all background processes.

    """

    def __init__(self, nprocesses = 16, progressMonitor = None):

        if nprocesses <= 0:
            raise ValueError("nprocesses must be at least one: " + str(nprocesses) + " is given")

        self.progressMonitor = NullProgressMonitor() if progressMonitor is None else progressMonitor
        self.nMaxProcesses = nprocesses
        self.nCurrentProcesses = 0
        self.task_queue = multiprocessing.JoinableQueue()
        self.result_queue = multiprocessing.Queue()
        self.lock = multiprocessing.Lock()
        self.nRunningTasks = 0
        self.taskNo = 0

    def begin(self):
        if self.nCurrentProcesses >= self.nMaxProcesses: return
        for i in xrange(self.nCurrentProcesses, self.nMaxProcesses):
            worker = Worker(
                task_queue = self.task_queue,
                result_queue = self.result_queue,
                progressReporter = self.progressMonitor.createReporter(),
                lock = self.lock
            )
            worker.start()
            self.nCurrentProcesses += 1

    def put(self, task, *args, **kwargs):
        self.task_queue.put((self.taskNo, task, args, kwargs))
        self.taskNo += 1
        self.nRunningTasks += 1

    def receive(self):
        messages = [ ] # a list of (taskNo, result)
        while self.nRunningTasks >= 1:
            if self.result_queue.empty(): continue
            message = self.result_queue.get()
            messages.append(message)
            self.nRunningTasks -= 1

        # sort in the order of taskNo
        messages = sorted(messages, key = itemgetter(0))

        results = [result for taskNo, result in messages]
        return results

    def end(self):
        for i in xrange(self.nCurrentProcesses):
            self.task_queue.put(None) # end workers
        self.task_queue.join()
        self.nCurrentProcesses = 0

##__________________________________________________________________||
<|endoftext|>


<commit_before>
import abc
from pox.openflow.libopenflow_01 import *
from pox.lib.packet.ethernet import *
import sts.headerspace.config_parser.openflow_parser as hsa

class Fingerprint(object):
  __metaclass__ = abc.ABCMeta

  # This should really be a protected constructor
  def __init__(self, field2value):
    # Make sure to convert arrays to tuples, since we need __hash__()
    for field, value in field2value.iteritems():
      if type(value) == list:
        field2value[field] = tuple(value)
    self._field2value = field2value

  def to_dict(self):
    return self._field2value

  @abc.abstractmethod
  def __hash__(self):
    pass

  @abc.abstractmethod
  def __eq__(self, other):
    pass

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

def process_data(msg):
  if msg.data == b'':
    return ()
  else:
    dp_packet = ethernet(msg.data)
    return DPFingerprint.from_pkt(dp_packet).to_dict()

def process_actions(msg):
  return tuple(map(str, map(type, msg.actions)))

class OFFingerprint(Fingerprint):
  ''' Fingerprints for openflow messages '''
  #  ofp_type -> fields to include in fingerprint
  # TODO(cs): I'm erring on the side of sparseness rather than completeness. We
  # may need to include more fields here to get an unambiguous fingerprint
  pkt_type_to_fields = {
    "ofp_features_reply" : ["datapath_id"],
    "ofp_switch_config" : ["flags"],
    "ofp_flow_mod" : ["match", "idle_timeout", "hard_timeout", "priority",
                     "out_port", "flags", "actions"],
    "ofp_port_mod" : ["port_no", "hw_addr", "config", "mask", "advertise"],
    "ofp_queue_get_config_request" : [],
    "ofp_queue_get_config_reply" : [],
    "ofp_stats_request" : ["type", "flags"],
    "ofp_stats_reply" : ["type", "flags"],
    "ofp_desc_stats" : [],
    "ofp_flow_stats_request" : [],
    "ofp_flow_stats" : [],
    "ofp_aggregate_stats_request" : [],
    "ofp_aggregate_stats" : [],
    "ofp_port_stats_request" : [],
    "ofp_port_stats" : [],
    "ofp_queue_stats_request" : [],
    "ofp_queue_stats" : [],
    "ofp_packet_out" : ["data", "in_port", "actions"],
    "ofp_barrier_reply" : [],
    "ofp_barrier_request" : [],
    "ofp_packet_in" : ["in_port", "data"],
    "ofp_flow_removed" : ["match", "reason", "priority"],
    "ofp_port_status" : ["reason", "desc"],
    "ofp_error" : ["type", "code"],
    "ofp_hello" : [],
    "ofp_echo_request" : [],
    "ofp_echo_reply" : [],
    "ofp_vendor_header" : [],
    #"ofp_vendor" (body of ofp_vendor_header)
    "ofp_features_request" : [],
    "ofp_get_config_request" : [],
    "ofp_get_config_reply" : [],
    "ofp_set_config" : []
  }
  special_fields = {
    # data needs a nested fingerprint
    'data' : process_data,
    # desc is a ofp_phy_port object
    'desc' : lambda pkt: (pkt.desc.port_no, pkt.desc.hw_addr.toStr()),
    # actions is an ordered list
    # for now, store it as a tuple of just the names of the action types
    'actions' : process_actions,
    # match has a bunch of crazy fields
    # Trick: convert it to an hsa match, and extract the human readable string
    # for the hsa match
    'match' : lambda pkt: hsa.hs_format["display"](hsa.ofp_match_to_hsa_match(pkt.match))
  }

  @staticmethod
  def from_pkt(pkt):
    pkt_type = type(pkt).__name__
    if pkt_type not in OFFingerprint.pkt_type_to_fields:
      raise ValueError("Unknown pkt_type %s" % pkt_type)
    field2value = {}
    field2value["class"] = pkt_type
    fields = OFFingerprint.pkt_type_to_fields[pkt_type]
    for field in fields:
      if field in OFFingerprint.special_fields:
        value = OFFingerprint.special_fields[field](pkt)
      else:
        value = getattr(pkt, field)
      field2value[field] = value
    return OFFingerprint(field2value)

  def __hash__(self):
    hash = 0
    class_name = self._field2value["class"]
    hash += class_name.__hash__()
    # Note that the order is important
    for field in self.pkt_type_to_fields[class_name]:
      hash += self._field2value[field].__hash__()
    return hash

  def __eq__(self, other):
    if type(other) != OFFingerprint:
      return False
    if self._field2value["class"] != other._field2value["class"]:
      return False
    klass = self._field2value["class"]
    for field in self.pkt_type_to_fields[klass]:
      if self._field2value[field] != other._field2value[field]:
        return False
    return True

class DPFingerprint(Fingerprint):
  ''' Fingerprints for dataplane messages '''
  fields = ['dl_src', 'dl_dst', 'nw_src', 'nw_dst']

  @staticmethod
  def from_pkt(pkt):
    # For now, just take (src MAC, dst MAC, src IP, dst IP) as the fingerprint for
    # dataplane packets
    # TODO(cs): might finer granularity later
    eth = pkt
    ip = pkt.next
    field2value = {'dl_src': eth.src.toStr(), 'dl_dst': eth.dst.toStr(),
                   'nw_src': ip.srcip.toStr(), 'nw_dst': ip.dstip.toStr()}
    return DPFingerprint(field2value)

  def __hash__(self):
    hash = 0
    # Note that the order is important
    for field in self.fields:
      hash += self._field2value[field].__hash__()
    return hash

  def __eq__(self, other):
    if type(other) != DPFingerprint:
      return False
    for field in self.fields:
      if self._field2value[field] != other._field2value[field]:
        return False
    return True
<commit_msg>Adding ofp_vendor to fingerprints<commit_after>
import abc
from pox.openflow.libopenflow_01 import *
from pox.lib.packet.ethernet import *
import sts.headerspace.config_parser.openflow_parser as hsa

class Fingerprint(object):
  __metaclass__ = abc.ABCMeta

  # This should really be a protected constructor
  def __init__(self, field2value):
    # Make sure to convert arrays to tuples, since we need __hash__()
    for field, value in field2value.iteritems():
      if type(value) == list:
        field2value[field] = tuple(value)
    self._field2value = field2value

  def to_dict(self):
    return self._field2value

  @abc.abstractmethod
  def __hash__(self):
    pass

  @abc.abstractmethod
  def __eq__(self, other):
    pass

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

def process_data(msg):
  if msg.data == b'':
    return ()
  else:
    dp_packet = ethernet(msg.data)
    return DPFingerprint.from_pkt(dp_packet).to_dict()

def process_actions(msg):
  return tuple(map(str, map(type, msg.actions)))

class OFFingerprint(Fingerprint):
  ''' Fingerprints for openflow messages '''
  #  ofp_type -> fields to include in fingerprint
  # TODO(cs): I'm erring on the side of sparseness rather than completeness. We
  # may need to include more fields here to get an unambiguous fingerprint
  pkt_type_to_fields = {
    "ofp_features_reply" : ["datapath_id"],
    "ofp_switch_config" : ["flags"],
    "ofp_flow_mod" : ["match", "idle_timeout", "hard_timeout", "priority",
                     "out_port", "flags", "actions"],
    "ofp_port_mod" : ["port_no", "hw_addr", "config", "mask", "advertise"],
    "ofp_queue_get_config_request" : [],
    "ofp_queue_get_config_reply" : [],
    "ofp_stats_request" : ["type", "flags"],
    "ofp_stats_reply" : ["type", "flags"],
    "ofp_desc_stats" : [],
    "ofp_flow_stats_request" : [],
    "ofp_flow_stats" : [],
    "ofp_aggregate_stats_request" : [],
    "ofp_aggregate_stats" : [],
    "ofp_port_stats_request" : [],
    "ofp_port_stats" : [],
    "ofp_queue_stats_request" : [],
    "ofp_queue_stats" : [],
    "ofp_packet_out" : ["data", "in_port", "actions"],
    "ofp_barrier_reply" : [],
    "ofp_barrier_request" : [],
    "ofp_packet_in" : ["in_port", "data"],
    "ofp_flow_removed" : ["match", "reason", "priority"],
    "ofp_port_status" : ["reason", "desc"],
    "ofp_error" : ["type", "code"],
    "ofp_hello" : [],
    "ofp_echo_request" : [],
    "ofp_echo_reply" : [],
    "ofp_vendor_header" : [],
    "ofp_vendor" : [], # (body of ofp_vendor_header)
    "ofp_features_request" : [],
    "ofp_get_config_request" : [],
    "ofp_get_config_reply" : [],
    "ofp_set_config" : []
  }
  special_fields = {
    # data needs a nested fingerprint
    'data' : process_data,
    # desc is a ofp_phy_port object
    'desc' : lambda pkt: (pkt.desc.port_no, pkt.desc.hw_addr.toStr()),
    # actions is an ordered list
    # for now, store it as a tuple of just the names of the action types
    'actions' : process_actions,
    # match has a bunch of crazy fields
    # Trick: convert it to an hsa match, and extract the human readable string
    # for the hsa match
    'match' : lambda pkt: hsa.hs_format["display"](hsa.ofp_match_to_hsa_match(pkt.match))
  }

  @staticmethod
  def from_pkt(pkt):
    pkt_type = type(pkt).__name__
    if pkt_type not in OFFingerprint.pkt_type_to_fields:
      raise ValueError("Unknown pkt_type %s" % pkt_type)
    field2value = {}
    field2value["class"] = pkt_type
    fields = OFFingerprint.pkt_type_to_fields[pkt_type]
    for field in fields:
      if field in OFFingerprint.special_fields:
        value = OFFingerprint.special_fields[field](pkt)
      else:
        value = getattr(pkt, field)
      field2value[field] = value
    return OFFingerprint(field2value)

  def __hash__(self):
    hash = 0
    class_name = self._field2value["class"]
    hash += class_name.__hash__()
    # Note that the order is important
    for field in self.pkt_type_to_fields[class_name]:
      hash += self._field2value[field].__hash__()
    return hash

  def __eq__(self, other):
    if type(other) != OFFingerprint:
      return False
    if self._field2value["class"] != other._field2value["class"]:
      return False
    klass = self._field2value["class"]
    for field in self.pkt_type_to_fields[klass]:
      if self._field2value[field] != other._field2value[field]:
        return False
    return True

class DPFingerprint(Fingerprint):
  ''' Fingerprints for dataplane messages '''
  fields = ['dl_src', 'dl_dst', 'nw_src', 'nw_dst']

  @staticmethod
  def from_pkt(pkt):
    # For now, just take (src MAC, dst MAC, src IP, dst IP) as the fingerprint for
    # dataplane packets
    # TODO(cs): might finer granularity later
    eth = pkt
    ip = pkt.next
    field2value = {'dl_src': eth.src.toStr(), 'dl_dst': eth.dst.toStr(),
                   'nw_src': ip.srcip.toStr(), 'nw_dst': ip.dstip.toStr()}
    return DPFingerprint(field2value)

  def __hash__(self):
    hash = 0
    # Note that the order is important
    for field in self.fields:
      hash += self._field2value[field].__hash__()
    return hash

  def __eq__(self, other):
    if type(other) != DPFingerprint:
      return False
    for field in self.fields:
      if self._field2value[field] != other._field2value[field]:
        return False
    return True
<|endoftext|>