<commit_before>from operator import attrgetter
import urllib.parse
import colander
from eliot import log_call
from ekklesia_portal.app import App
from ekklesia_portal.concepts.argument_relation.argument_relations import ArgumentRelations
from ekklesia_portal.concepts.ekklesia_portal.cell.layout import LayoutCell
from ekklesia_portal.concepts.ekklesia_portal.cell.form import NewFormCell
from ekklesia_portal.concepts.ekklesia_portal.cell.form import EditFormCell
from ekklesia_portal.database.datamodel import Department, Proposition, Tag, PropositionNote
from ekklesia_common.translation import _
from ekklesia_common.cell import Cell
from ekklesia_portal.enums import ArgumentType, PropositionStatus, OpenSlidesVotingResult
from ekklesia_portal.permission import SupportPermission, CreatePermission, EditPermission
from .propositions import Propositions
from .proposition_helper import items_for_proposition_select_widgets

@App.cell(Proposition)
class PropositionCell(LayoutCell):

    model_properties = [
        'abstract',
        'ballot',
        'content',
        'created_at',
        'derivations',
        'external_discussion_url',
        'id',
        'modifies',
        'motivation',
        'replacements',
        'replaces',
        'tags',
        'title'
    ]

    actions = Cell.fragment('proposition_actions')
    tabs = Cell.fragment('proposition_tabs')
    small = Cell.fragment('proposition_small')
    card = Cell.fragment('proposition_card')

    @Cell.fragment
    def status(self):
        status_to_variant = {
            PropositionStatus.DRAFT: 'submitted',
            PropositionStatus.CHANGING: 'submitted',
            PropositionStatus.SUBMITTED: 'submitted',
            PropositionStatus.ABANDONED: 'submitted',
            PropositionStatus.QUALIFIED: 'submitted',
            PropositionStatus.SCHEDULED: 'scheduled',
            PropositionStatus.VOTING: 'scheduled',
            PropositionStatus.FINISHED: 'finished',
        }
        variant = status_to_variant[self._model.status]
        template = f"proposition/proposition_status_{variant}.j2.jade"
        return self.render_template(template)

    @Cell.fragment
    def history(self):
        status_to_variant = {
            PropositionStatus.DRAFT: 'submitted',
            PropositionStatus.CHANGING: 'submitted',
            PropositionStatus.SUBMITTED: 'submitted',
            PropositionStatus.ABANDONED: 'submitted',
            PropositionStatus.QUALIFIED: 'submitted',
            PropositionStatus.SCHEDULED: 'scheduled',
            PropositionStatus.VOTING: 'scheduled',
            PropositionStatus.FINISHED: 'finished',
        }
        variant = status_to_variant[self._model.status]
        template = f"proposition/proposition_history_{variant}.j2.jade"
        return self.render_template(template)

    def associated_url(self):
        return self.link(self._model, 'associated')

    def report_url(self):
        return self._s.app.report_url

    def share_url(self):
        if self._app.settings.share.use_url_shortener:
            from ekklesia_portal.helper.url_shortener import make_tiny
            return make_tiny(self.self_link)
        else:
            return self.self_link[:69]

    @log_call
    def share_email_url(self):
        share_email_topic = (self._s.share.email_topic[self.language]
                             .format(
                                 voting_identifier=self._model.voting_identifier,
                                 title=self._model.title[:140]))

        share_email_body = self._s.share.email_body[self.language] + self.share_url
        email_url = urllib.parse.urlencode({'subject': share_email_topic,
                                            'body': share_email_body},
                                           quote_via=urllib.parse.quote)
        email_url = 'mailto:?' + email_url
        return email_url

    @log_call
    def share_twitter_url(self):
        share_tweet_msg = (self._s.share.tweet_msg[self.language]
                           .format(
                               voting_identifier=self._model.voting_identifier,
                               title=self._model.title[:70]))

        twitter_url = urllib.parse.urlencode({'hashtags': self._app.settings.share.hashtag,
                                              'related': self._app.settings.share.promote_account,
                                              'text': share_tweet_msg,
                                              'tw_p': 'tweetbutton',
                                              'url': self.share_url})
        twitter_url = 'https://twitter.com/intent/tweet?' + twitter_url
        return twitter_url

    def ballot_url(self):
        return self.link(self._model.ballot)

    def ballot_title(self):
        ballot = self._model.ballot
        if ballot.name:
            return ballot.name

        return f'# {ballot.id}'

    def discussion_url(self):
        return self.link(self._model)

    def propositions_tag_url(self, tag):
        return self.class_link(Propositions, dict(tag=tag.name))

    def is_supported_by_current_user(self):
        if self.current_user is None:
            return False
        return self._model.support_by_user(self.current_user) is not None

    def discussion_link_class(self):
        return 'active' if self.options.get('active_tab') == 'discussion' else ''

    def associated_link_class(self):
        return 'active' if self.options.get('active_tab') == 'associated' else ''

    def new_associated_proposition_url(self, association_type):
        return self.class_link(Propositions, dict(association_type=association_type), '+new')

    def new_pro_argument_url(self):
        return self.class_link(
            ArgumentRelations,
            dict(proposition_id=self._model.id, relation_type=ArgumentType.PRO.name),
            '+new')

    def new_con_argument_url(self):
        return self.class_link(
            ArgumentRelations,
            dict(proposition_id=self._model.id, relation_type=ArgumentType.CONTRA.name),
            '+new')

    def supporter_count(self):
        return self._model.active_supporter_count

    def support_action(self):
        return self.link(self._model, 'support')

    def pro_argument_relations(self):
        relations = [p for p in self._model.proposition_arguments if p.argument_type == ArgumentType.PRO]
        return sorted(relations, key=attrgetter('score'), reverse=True)

    def contra_argument_relations(self):
        relations = [p for p in self._model.proposition_arguments if p.argument_type == ArgumentType.CONTRA]
        return sorted(relations, key=attrgetter('score'), reverse=True)

    def argument_count(self):
        return len(self._model.proposition_arguments)

    def full_title(self):
        if self._model.voting_identifier:
            return self._model.voting_identifier + ': ' + self._model.title
        else:
            return self._model.title

    def show_support_actions(self):
        return self._request.permitted_for_current_user(self._model, SupportPermission)

    def show_create_argument(self):
        return self._request.permitted_for_current_user(ArgumentRelations(), CreatePermission)

    def show_create_associated_proposition(self):
        return self._request.permitted_for_current_user(self._model, CreatePermission)

    def voting_phase(self):
        return self._model.ballot.voting

    def voting_result_state(self):
        result = self._model.ballot.result
        if result:
            try:
                return OpenSlidesVotingResult(result.get(self._model.voting_identifier, {}).get("state"))
            except ValueError:
                return

    def voting_result_symbol(self):
        symbols = {
            OpenSlidesVotingResult.ACCEPTED: "fas fa-check",
            OpenSlidesVotingResult.REJECTED: "fas fa-ban",
            OpenSlidesVotingResult.NOT_DECIDED: "fas fa-spinner"
        }
        return symbols.get(self.voting_result_state)

    def show_edit_button(self):
        return self.options.get('show_edit_button') and self._request.permitted_for_current_user(self._model, EditPermission)

    def edit_url(self):
        return self.link(self._model, 'edit')

    def note_url(self):
        if self.current_user is None:
            return False
        return self.class_link(
            PropositionNote,
            variables={'proposition_id': self._model.id, 'user_id': self._request.current_user.id},
            name='edit')

@App.cell(Propositions, 'new')
class NewPropositionCell(NewFormCell):

    def _prepare_form_for_render(self):
        if self._request.identity.has_global_admin_permissions:
            departments = self._request.q(Department)
        else:
            departments = self._request.current_user.departments

        tags = self._request.q(Tag).all()

        if self._form.error is None or self._form.cstruct['tags'] is colander.null:
            selected_tags = None
        else:
            selected_tags = self._form.cstruct['tags'] if self._form.error is not None else None

        items = items_for_proposition_select_widgets(departments, tags, selected_tags)
        self._form.prepare_for_render(items)

class EditPropositionCell(EditFormCell):

    def _prepare_form_for_render(self):
        form_data = self._model.to_dict()
        form_data['area_id'] = self._model.ballot.area_id
        self.set_form_data(form_data)

        if self._request.identity.has_global_admin_permissions:
            departments = self._request.q(Department)
        else:
            departments = self._request.current_user.departments

        tags = self._request.q(Tag).all()
        items = items_for_proposition_select_widgets(departments, tags)
        self._form.prepare_for_render(items)

    def department_name(self):
        return self._model.ballot.area.department.name

    def subject_area_name(self):
        return self._model.ballot.area.name

    def ballot_id(self):
        return self._model.ballot.id

    def ballot_name(self):
        return self._model.ballot.name

    def ballot_url(self):
        return self.link(self._model.ballot)

@App.cell(Propositions)
class PropositionsCell(LayoutCell):

    model_properties = ['mode', 'tag', 'search']

    def propositions(self):
        return list(self._model.propositions(self._request.q))
<commit_msg>Proposition, edit: render entered values on validation error, fix tags<commit_after>from operator import attrgetter
import urllib.parse
import colander
from eliot import log_call
from ekklesia_portal.app import App
from ekklesia_portal.concepts.argument_relation.argument_relations import ArgumentRelations
from ekklesia_portal.concepts.ekklesia_portal.cell.layout import LayoutCell
from ekklesia_portal.concepts.ekklesia_portal.cell.form import NewFormCell
from ekklesia_portal.concepts.ekklesia_portal.cell.form import EditFormCell
from ekklesia_portal.database.datamodel import Department, Proposition, Tag, PropositionNote
from ekklesia_common.translation import _
from ekklesia_common.cell import Cell
from ekklesia_portal.enums import ArgumentType, PropositionStatus, OpenSlidesVotingResult
from ekklesia_portal.permission import SupportPermission, CreatePermission, EditPermission
from .propositions import Propositions
from .proposition_helper import items_for_proposition_select_widgets

@App.cell(Proposition)
class PropositionCell(LayoutCell):

    model_properties = [
        'abstract',
        'ballot',
        'content',
        'created_at',
        'derivations',
        'external_discussion_url',
        'id',
        'modifies',
        'motivation',
        'replacements',
        'replaces',
        'tags',
        'title'
    ]

    actions = Cell.fragment('proposition_actions')
    tabs = Cell.fragment('proposition_tabs')
    small = Cell.fragment('proposition_small')
    card = Cell.fragment('proposition_card')

    @Cell.fragment
    def status(self):
        status_to_variant = {
            PropositionStatus.DRAFT: 'submitted',
            PropositionStatus.CHANGING: 'submitted',
            PropositionStatus.SUBMITTED: 'submitted',
            PropositionStatus.ABANDONED: 'submitted',
            PropositionStatus.QUALIFIED: 'submitted',
            PropositionStatus.SCHEDULED: 'scheduled',
            PropositionStatus.VOTING: 'scheduled',
            PropositionStatus.FINISHED: 'finished',
        }
        variant = status_to_variant[self._model.status]
        template = f"proposition/proposition_status_{variant}.j2.jade"
        return self.render_template(template)

    @Cell.fragment
    def history(self):
        status_to_variant = {
            PropositionStatus.DRAFT: 'submitted',
            PropositionStatus.CHANGING: 'submitted',
            PropositionStatus.SUBMITTED: 'submitted',
            PropositionStatus.ABANDONED: 'submitted',
            PropositionStatus.QUALIFIED: 'submitted',
            PropositionStatus.SCHEDULED: 'scheduled',
            PropositionStatus.VOTING: 'scheduled',
            PropositionStatus.FINISHED: 'finished',
        }
        variant = status_to_variant[self._model.status]
        template = f"proposition/proposition_history_{variant}.j2.jade"
        return self.render_template(template)

    def associated_url(self):
        return self.link(self._model, 'associated')

    def report_url(self):
        return self._s.app.report_url

    def share_url(self):
        if self._app.settings.share.use_url_shortener:
            from ekklesia_portal.helper.url_shortener import make_tiny
            return make_tiny(self.self_link)
        else:
            return self.self_link[:69]

    @log_call
    def share_email_url(self):
        share_email_topic = (self._s.share.email_topic[self.language]
                             .format(
                                 voting_identifier=self._model.voting_identifier,
                                 title=self._model.title[:140]))

        share_email_body = self._s.share.email_body[self.language] + self.share_url
        email_url = urllib.parse.urlencode({'subject': share_email_topic,
                                            'body': share_email_body},
                                           quote_via=urllib.parse.quote)
        email_url = 'mailto:?' + email_url
        return email_url

    @log_call
    def share_twitter_url(self):
        share_tweet_msg = (self._s.share.tweet_msg[self.language]
                           .format(
                               voting_identifier=self._model.voting_identifier,
                               title=self._model.title[:70]))

        twitter_url = urllib.parse.urlencode({'hashtags': self._app.settings.share.hashtag,
                                              'related': self._app.settings.share.promote_account,
                                              'text': share_tweet_msg,
                                              'tw_p': 'tweetbutton',
                                              'url': self.share_url})
        twitter_url = 'https://twitter.com/intent/tweet?' + twitter_url
        return twitter_url

    def ballot_url(self):
        return self.link(self._model.ballot)

    def ballot_title(self):
        ballot = self._model.ballot
        if ballot.name:
            return ballot.name

        return f'# {ballot.id}'

    def discussion_url(self):
        return self.link(self._model)

    def propositions_tag_url(self, tag):
        return self.class_link(Propositions, dict(tag=tag.name))

    def is_supported_by_current_user(self):
        if self.current_user is None:
            return False
        return self._model.support_by_user(self.current_user) is not None

    def discussion_link_class(self):
        return 'active' if self.options.get('active_tab') == 'discussion' else ''

    def associated_link_class(self):
        return 'active' if self.options.get('active_tab') == 'associated' else ''

    def new_associated_proposition_url(self, association_type):
        return self.class_link(Propositions, dict(association_type=association_type), '+new')

    def new_pro_argument_url(self):
        return self.class_link(
            ArgumentRelations,
            dict(proposition_id=self._model.id, relation_type=ArgumentType.PRO.name),
            '+new')

    def new_con_argument_url(self):
        return self.class_link(
            ArgumentRelations,
            dict(proposition_id=self._model.id, relation_type=ArgumentType.CONTRA.name),
            '+new')

    def supporter_count(self):
        return self._model.active_supporter_count

    def support_action(self):
        return self.link(self._model, 'support')

    def pro_argument_relations(self):
        relations = [p for p in self._model.proposition_arguments if p.argument_type == ArgumentType.PRO]
        return sorted(relations, key=attrgetter('score'), reverse=True)

    def contra_argument_relations(self):
        relations = [p for p in self._model.proposition_arguments if p.argument_type == ArgumentType.CONTRA]
        return sorted(relations, key=attrgetter('score'), reverse=True)

    def argument_count(self):
        return len(self._model.proposition_arguments)

    def full_title(self):
        if self._model.voting_identifier:
            return self._model.voting_identifier + ': ' + self._model.title
        else:
            return self._model.title

    def show_support_actions(self):
        return self._request.permitted_for_current_user(self._model, SupportPermission)

    def show_create_argument(self):
        return self._request.permitted_for_current_user(ArgumentRelations(), CreatePermission)

    def show_create_associated_proposition(self):
        return self._request.permitted_for_current_user(self._model, CreatePermission)

    def voting_phase(self):
        return self._model.ballot.voting

    def voting_result_state(self):
        result = self._model.ballot.result
        if result:
            try:
                return OpenSlidesVotingResult(result.get(self._model.voting_identifier, {}).get("state"))
            except ValueError:
                return

    def voting_result_symbol(self):
        symbols = {
            OpenSlidesVotingResult.ACCEPTED: "fas fa-check",
            OpenSlidesVotingResult.REJECTED: "fas fa-ban",
            OpenSlidesVotingResult.NOT_DECIDED: "fas fa-spinner"
        }
        return symbols.get(self.voting_result_state)

    def show_edit_button(self):
        return self.options.get('show_edit_button') and self._request.permitted_for_current_user(self._model, EditPermission)

    def edit_url(self):
        return self.link(self._model, 'edit')

    def note_url(self):
        if self.current_user is None:
            return False
        return self.class_link(
            PropositionNote,
            variables={'proposition_id': self._model.id, 'user_id': self._request.current_user.id},
            name='edit')

@App.cell(Propositions, 'new')
class NewPropositionCell(NewFormCell):

    def _prepare_form_for_render(self):
        if self._request.identity.has_global_admin_permissions:
            departments = self._request.q(Department)
        else:
            departments = self._request.current_user.departments

        tags = self._request.q(Tag).all()

        if self._form.error is None or self._form.cstruct['tags'] is colander.null:
            selected_tag_names = None
        else:
            selected_tag_names = self._form.cstruct['tags']

        items = items_for_proposition_select_widgets(departments, tags, selected_tag_names)
        self._form.prepare_for_render(items)

class EditPropositionCell(EditFormCell):

    def _prepare_form_for_render(self):
        if self._form.error is None:
            form_data = self._model.to_dict()
            form_data['area_id'] = self._model.ballot.area_id
            form_data['tags'] = [t.name for t in self._model.tags]
            self.set_form_data(form_data)
            selected_tag_names = None
        else:
            selected_tag_names = self._form.cstruct['tags']

        if self._request.identity.has_global_admin_permissions:
            departments = self._request.q(Department)
        else:
            departments = self._request.current_user.departments

        tags = self._request.q(Tag).all()
        items = items_for_proposition_select_widgets(departments, tags, selected_tag_names)
        self._form.prepare_for_render(items)

    def department_name(self):
        return self._model.ballot.area.department.name

    def subject_area_name(self):
        return self._model.ballot.area.name

    def ballot_id(self):
        return self._model.ballot.id

    def ballot_name(self):
        return self._model.ballot.name

    def ballot_url(self):
        return self.link(self._model.ballot)

@App.cell(Propositions)
class PropositionsCell(LayoutCell):

    model_properties = ['mode', 'tag', 'search']

    def propositions(self):
        return list(self._model.propositions(self._request.q))
<|endoftext|>


<commit_before>#!/usr/bin/env python
"""

Use Monte Carlo to estimate the probability of outliving
one's savings in retirement.

Uses historical economic data to estimate investment return,
and CDC life tables to estimate chance of death.

"""

import numpy as np
import pandas as pd
import uncertainties as unc
import uncertainties.unumpy as unp

import matplotlib.pyplot as plt
from matplotlib import rcParams

from itertools import cycle
lines = ["-","--","-.",":"]
linecycler = cycle(lines)

from cdc_life_tables import life_table
import shiller

# Historical financial data
inflation = shiller.inflation[1:-1]
stock_returns = shiller.stock_returns[1:-1]
interest_rates = shiller.interest_rates[1:-1]

rand = np.random.random_sample

def run_histories(starting_assets,
                  yearly_expenses,
                  stock_fraction,
                  starting_age,
                  state_abbrev,
                  demographic_group,
                  n_mc=1000, plotting=False, verbose=False):
    """
    Run a Monte Carlo simulation for a person starting with the given
    amount of assets in savings. The yearly_expenses are withdrawn
    once a year, while the assets grow according to historical
    US stock and bond returns, dividing between the two according to
    stock_fraction. The chance of dying each year is taken from
    life tables.

    Sources of data:

       * Shiller's historical economic data
       * CDC's life tables

    Inputs:

       * starting_assets : amount of initial savings to invest for income
       * yearly_expenses : amount of money needed per year. This value
                            will be adjusted for inflation.
       * stock_fraction : fraction (between 0.0 and 1.0) of the money invested
                            in stocks. The remainder is invested in bonds.
       * starting_age : the subject's age at which yearly withdraws will be made from
                            the investment
       * state_abbrev : mailing abbreviation for the state in which the subject lives
       * demographic_group : the subject's demographic group accepted by
                               cdc_life_tables.life_table
       * n_mc : the number of Monte Carlo histories
       * plotting : produce a plot showing the Monte Carlo histories
       * verbose : produce verbose diagnostic messages

    Output:
       * probability of running out of money

    """

    # Life table
    table = life_table(state_abbrev, demographic_group)

    mc_histories = []

    for i in range(n_mc):

        age = starting_age
        current_assets = starting_assets
        expenses_per_year = yearly_expenses

        assets = [current_assets]

        # Loop over years
        while current_assets > 0:

            # Death this year.
            if age >= 110 or rand() <= table[int(age)]:
                # Die at random point in year
                current_assets -= expenses_per_year*rand()
                break

            # Subtracting expenses for year
            current_assets -= expenses_per_year

            # Pick past year by random to base inflation, stock return data
            i = np.random.randint(inflation.size, size=1)
            i = int(i)

            # Adjust expenses for inflation.
            expenses_per_year *= 1.0+inflation[i]

            # Adding stock investment increase
            stock_gains = stock_returns[i] * (current_assets*stock_fraction)

            # Adding bond investment increase
            bond_gains = interest_rates[i] * (current_assets*(1-stock_fraction))

            current_assets += stock_gains
            current_assets += bond_gains

            # Saving current assets
            assets.append(current_assets)

            # Getting old
            age += 1.0

        assets = np.array(assets)

        mc_histories.append( (assets) )

    if plotting:

        rcParams['figure.figsize'] = [7.0, 3.5]

        # Plot of asset-over-lifetime histories
        plt.figure()

        final_ages = []
        final_assets = []
        for i in range(n_mc):
            y = mc_histories[i] / 1e6
            x = np.arange(starting_age, starting_age+y.size)
            plt.plot(x, y, color='gray', linewidth=0.5)

            final_assets.append(y[-1])
            final_ages.append(x[-1])

        plt.plot(final_ages, final_assets, color='red', ls='.',
                 marker='.', markersize=1.5)

        plt.xlabel('Age')
        plt.ylabel('Remaining assets (million USD)')

        # plt.savefig('figs/histories.pdf')

        # Plot of age of death
        rcParams['figure.subplot.left'] = 0.15
        plt.figure()

        final_ages = np.array(final_ages)

        max_age = 110
        bins = np.linspace(int(starting_age)-0.5, max_age+0.5, max_age-starting_age+2)

        plt.hist(final_ages, bins=bins, normed=True)

        plt.xlabel('Age of Death')
        plt.ylabel('Probability')

        # plt.savefig('figs/final-age.pdf')

    else:
        final_assets = []
        for i in range(n_mc):
            final_assets.append(mc_histories[i][-1])

    final_assets = np.array(final_assets)

    run_out_of_money_hist = np.array(final_assets < 0.0, dtype=np.float64)
    run_out_of_money = unc.ufloat(run_out_of_money_hist.mean(),
                                  run_out_of_money_hist.std()/np.sqrt(n_mc))

    if verbose:
        print ' Chance of running out of money is {:%}'.format(run_out_of_money)

    return run_out_of_money

def cascade_plot(yearly_expenses,
                 stock_fraction,
                 starting_age,
                 state_abbrev,
                 demographic_group,
                 stock_fractions = [0.25, 0.5, 0.75],
                 n_mc=5000):
    """
    Inputs:

       * yearly_expenses : amount of money needed per year. This value
                            will be adjusted for inflation.
       * stock_fraction : fraction (between 0.0 and 1.0) of the money invested
                            in stocks. The remainder is invested in bonds.
       * starting_age : the subject's age at which yearly withdraws will be made from
                            the investment
       * state_abbrev : mailing abbreviation for the state in which the subject lives
       * demographic_group : the subject's demographic group accepted by
                               cdc_life_tables.life_table
       * n_mc : the number of Monte Carlo histories
       * plotting : produce a plot showing the Monte Carlo histories
       * verbose : produce verbose diagnostic messages

    Output:
       * Matplotlib figure object

    """
    fig = plt.figure()

    starting_assets = np.linspace(1e5, 10e6, 100)
    starting_assets = np.array(starting_assets)

    for stock_fraction in stock_fractions:

        run_out_of_money = []
        for x in starting_assets:
            p = 100*run_histories(x,
                                  yearly_expenses,
                                  stock_fraction,
                                  starting_age,
                                  state_abbrev,
                                  demographic_group,
                                  n_mc=n_mc)
            run_out_of_money.append(p)

            # Don't keep going for probability <1%
            if p < 1: break

        run_out_of_money = np.array(run_out_of_money)
        n = run_out_of_money.size

        plt.errorbar(starting_assets[:n]/1e6, unp.nominal_values(run_out_of_money),
                     yerr=unp.std_devs(run_out_of_money),
                     capsize=0.0, marker='.', markersize=3.5, ls=next(linecycler),
                     label='{:.0%} stocks'.format(stock_fraction))

    plt.xlabel('Starting Assets (million USD)')
    plt.ylabel('Prob. of running out of money (%)')

    str_id = '{}-{}-{}-{}'.format(demographic_group, state_abbrev, starting_age,
                                                yearly_expenses)
    plt.title('{}-{}, starting at age {} with \${}/year expenses'.format(demographic_group, state_abbrev, starting_age,
                                                yearly_expenses))

    plt.legend(fontsize='x-small')
    plt.ylim(ymin=0, ymax=100)

    # plt.savefig('figs/{}.pdf'.format(str_id))

    return fig

if __name__ == '__main__':
    starting_age = 65.0
    state_abbrev = 'IA'
    demographic_group = 'mf'

    # Expenses per year
    yearly_expenses = 50e3

    # Assets
    starting_assets = 3e6

    # Investment allocation
    stock_fraction = 0.5

    # Run one simulation for a given starting assets and stock fraction
    run_histories(starting_assets, yearly_expenses,
                  stock_fraction,
                  starting_age,
                  state_abbrev,
                  demographic_group,
                  n_mc=1000, plotting=True, verbose=True)

    # Run many simulations over a range of starting assets and stock fractions
    cascade_plot(yearly_expenses,
                 stock_fraction,
                 starting_age,
                 state_abbrev,
                 demographic_group)
<commit_msg>Created sensitivity plots for each independent variable<commit_after>#!/usr/bin/env python
"""

Use Monte Carlo to estimate the probability of outliving
one's savings in retirement.

Uses historical economic data to estimate investment return,
and CDC life tables to estimate chance of death.

"""

import numpy as np
from scipy.optimize import brentq
import pandas as pd
import uncertainties as unc
import uncertainties.unumpy as unp

import matplotlib.pyplot as plt
from matplotlib import rcParams

from itertools import cycle
lines = ["-","--","-.",":"]
linecycler = cycle(lines)

from cdc_life_tables import life_table
import shiller

# Historical financial data
inflation = shiller.inflation.iloc[1:-1]
stock_returns = shiller.stock_returns.iloc[1:-1]
interest_rates = shiller.interest_rates.iloc[1:-1]

rand = np.random.random_sample

def run_histories(starting_assets,
                  yearly_expense,
                  stock_fraction,
                  starting_age,
                  state_abbrev,
                  demographic_group,
                  n_mc=1000, plotting=False, verbose=False):
    """
    Run a Monte Carlo simulation for a person starting with the given
    amount of assets in savings. The yearly_expense are withdrawn
    once a year, while the assets grow according to historical
    US stock and bond returns, dividing between the two according to
    stock_fraction. The chance of dying each year is taken from
    life tables.

    Sources of data:

       * Shiller's historical economic data
       * CDC's life tables

    Inputs:

       * starting_assets : amount of initial savings to invest for income
       * yearly_expense  : amount of money needed per year. This value
                            will be adjusted for inflation.
       * stock_fraction : fraction (between 0.0 and 1.0) of the money invested
                            in stocks. The remainder is invested in bonds.
       * starting_age : the subject's age at which yearly withdraws will be made from
                            the investment
       * state_abbrev : mailing abbreviation for the state in which the subject lives
       * demographic_group : the subject's demographic group accepted by
                               cdc_life_tables.life_table
       * n_mc : the number of Monte Carlo histories
       * plotting : produce a plot showing the Monte Carlo histories
       * verbose : produce verbose diagnostic messages

    Output:
       * probability of running out of money

    """

    # Life table
    table = life_table(state_abbrev, demographic_group)

    mc_histories = []

    for i in range(n_mc):

        age = starting_age
        current_assets = starting_assets
        expenses_per_year = yearly_expense

        assets = [current_assets]

        # Loop over years
        while current_assets > 0:

            # Death this year.
            if age >= 110 or rand() <= table[int(age)]:
                # Die at random point in year
                current_assets -= expenses_per_year*rand()
                break

            # Subtracting expenses for year
            current_assets -= expenses_per_year

            # Pick past year by random to base inflation, stock return data
            i = np.random.randint(inflation.size, size=1)
            i = int(i)

            # Adjust expenses for inflation.
            expenses_per_year *= 1.0+inflation.iloc[i]

            # Adding stock investment increase
            stock_gains = stock_returns.iloc[i] * (current_assets*stock_fraction)

            # Adding bond investment increase
            bond_gains = interest_rates.iloc[i] * (current_assets*(1-stock_fraction))

            current_assets += stock_gains
            current_assets += bond_gains

            # Saving current assets
            assets.append(current_assets)

            # Getting old
            age += 1.0

        assets = np.array(assets)

        mc_histories.append( (assets) )

    if plotting:

        rcParams['figure.figsize'] = [7.0, 3.5]

        # Plot of asset-over-lifetime histories
        plt.figure()

        final_ages = []
        final_assets = []
        for i in range(n_mc):
            y = mc_histories[i] / 1e6
            x = np.arange(starting_age, starting_age+y.size)
            plt.plot(x, y, color='gray', linewidth=0.5)

            final_assets.append(y[-1])
            final_ages.append(x[-1])

        plt.plot(final_ages, final_assets, color='red', ls='.',
                 marker='.', markersize=1.5)

        plt.xlabel('Age')
        plt.ylabel('Remaining assets (million USD)')

        # plt.savefig('figs/histories.pdf')

        # Plot of age of death
        rcParams['figure.subplot.left'] = 0.15
        plt.figure()

        final_ages = np.array(final_ages)

        max_age = 110
        bins = np.linspace(int(starting_age)-0.5, max_age+0.5, max_age-starting_age+2)

        plt.hist(final_ages, bins=bins, normed=True)

        plt.xlabel('Age of Death')
        plt.ylabel('Probability')

        # plt.savefig('figs/final-age.pdf')

    else:
        final_assets = []
        for i in range(n_mc):
            final_assets.append(mc_histories[i][-1])

    final_assets = np.array(final_assets)

    run_out_of_money_hist = np.array(final_assets < 0.0, dtype=np.float64)
    run_out_of_money = unc.ufloat(run_out_of_money_hist.mean(),
                                  run_out_of_money_hist.std()/np.sqrt(n_mc))

    if verbose:
        print ' Chance of running out of money is {:%}'.format(run_out_of_money)

    return run_out_of_money

def how_much_to_save(
                     acceptable_risk=0.01,
                     yearly_expense=40e3,
                     stock_fraction=0.5,
                     starting_age=65,
                     state_abbrev='CA',
                     demographic_group='total',
                     n_mc=1000, plotting=False, verbose=False):
    """
    Computes f(x) = f_0, where f is the MC simulation of the retirement
    process returning the probability of running out of money and
    x is the size of the starting assets.

    Inputs:

       * yearly_expense : amount of money needed per year. This value
                            will be adjusted for inflation.
       * stock_fraction : fraction (between 0.0 and 1.0) of the money invested
                            in stocks. The remainder is invested in bonds.
       * starting_age : the subject's age at which yearly withdraws will be made from
                            the investment
       * state_abbrev : mailing abbreviation for the state in which the subject lives
       * demographic_group : the subject's demographic group accepted by
                               cdc_life_tables.life_table
       * acceptable_risk : probability of running out of money
       * n_mc : the number of Monte Carlo histories
       * plotting : produce a plot showing the Monte Carlo histories
       * verbose : produce verbose diagnostic messages

    Output:

       * starting_assets : amount of initial savings to invest for income

    """

    # Probabilities away from 0 and 1 have larger statistical noise
    #  and so need more histories.
    #if 0.3 < acceptable_risk < 0.7:
    #    n_mc = 1000
    #elif 0.1 < acceptable_risk < 0.9:
    #    n_mc = 800
    #else:
    #    n_mc = 400
    n_mc = 500

    def f(x):
        prob_outlive_savings = run_histories(x, yearly_expense, stock_fraction,
                                             starting_age, state_abbrev,
                                             demographic_group,
                                             n_mc=n_mc, plotting=False, verbose=False)
        return acceptable_risk - prob_outlive_savings.nominal_value

    lo_bound = 3.0*yearly_expense
    hi_bound = 200.0*yearly_expense
    res = brentq(f, lo_bound, hi_bound, rtol=1e-2, full_output=True)

    return res[0]

def cascade_plot(yearly_expense,
                 stock_fraction,
                 starting_age,
                 state_abbrev,
                 demographic_group,
                 stock_fractions = [0.25, 0.5, 0.75],
                 n_mc=5000):
    """
    Inputs:

       * yearly_expense : amount of money needed per year. This value
                            will be adjusted for inflation.
       * stock_fraction : fraction (between 0.0 and 1.0) of the money invested
                            in stocks. The remainder is invested in bonds.
       * starting_age : the subject's age at which yearly withdraws will be made from
                            the investment
       * state_abbrev : mailing abbreviation for the state in which the subject lives
       * demographic_group : the subject's demographic group accepted by
                               cdc_life_tables.life_table
       * n_mc : the number of Monte Carlo histories
       * plotting : produce a plot showing the Monte Carlo histories
       * verbose : produce verbose diagnostic messages

    Output:
       * Matplotlib figure object

    """
    fig = plt.figure()

    starting_assets = np.linspace(1e5, 10e6, 100)
    starting_assets = np.array(starting_assets)

    for stock_fraction in stock_fractions:

        run_out_of_money = []
        for x in starting_assets:
            p = 100*run_histories(x,
                                  yearly_expense,
                                  stock_fraction,
                                  starting_age,
                                  state_abbrev,
                                  demographic_group,
                                  n_mc=n_mc)
            run_out_of_money.append(p)

            # Don't keep going for probability <1%
            if p < 1: break

        run_out_of_money = np.array(run_out_of_money)
        n = run_out_of_money.size

        plt.errorbar(starting_assets[:n]/1e6, unp.nominal_values(run_out_of_money),
                     yerr=unp.std_devs(run_out_of_money),
                     capsize=0.0, marker='.', markersize=3.5, ls=next(linecycler),
                     label='{:.0%} stocks'.format(stock_fraction))

    plt.xlabel('Starting Assets (million USD)')
    plt.ylabel('Prob. of running out of money (%)')

    str_id = '{}-{}-{}-{}'.format(demographic_group, state_abbrev, starting_age,
                                                yearly_expense)
    plt.title('{}-{}, starting at age {} with \${}/year expenses'.format(demographic_group, state_abbrev, starting_age,
                                                yearly_expense))

    plt.legend(fontsize='x-small')
    plt.ylim(ymin=0, ymax=100)

    # plt.savefig('figs/{}.pdf'.format(str_id))

    return fig

def sensitivity_plots(
                      state_abbrev='CA',
                      demographic_group='total',
                      yearly_expense=40e3,
                      yearly_expenses=np.logspace(3.69897, 5, 10),
                      starting_age=65,
                      starting_ages=np.linspace(40, 85, 10),
                      acceptable_risk=0.02,
                      acceptable_risks=np.logspace(-3, -0.2, 7),
                      stock_fraction=0.5,
                      stock_fractions=np.linspace(0.0, 1.0, 11),
                      n_mc=5000):
    """
    Inputs:

       * yearly_expense : amount of money needed per year. This value
                            will be adjusted for inflation.
       * stock_fraction : fraction (between 0.0 and 1.0) of the money invested
                            in stocks. The remainder is invested in bonds.
       * starting_age : the subject's age at which yearly withdraws will be made from
                            the investment
       * state_abbrev : mailing abbreviation for the state in which the subject lives
       * demographic_group : the subject's demographic group accepted by
                               cdc_life_tables.life_table
       * n_mc : the number of Monte Carlo histories
       * plotting : produce a plot showing the Monte Carlo histories
       * verbose : produce verbose diagnostic messages

    Output:
       * Matplotlib figure object

    """

    factors = {
                'stock_fraction'  : {'value' : stock_fraction,  'values' : stock_fractions },
                'acceptable_risk' : {'value' : acceptable_risk, 'values' : acceptable_risks },
                 'yearly_expense' : {'value' : yearly_expense,  'values' : yearly_expenses },
                   'starting_age' : {'value' : starting_age,    'values' : starting_ages },
              }

    rcParams['figure.figsize'] = [9, 11]
    fig, axs = plt.subplots(nrows=len(factors.keys()))

    base_opts = {
                  'stock_fraction'    : stock_fraction,
                  'acceptable_risk'   : acceptable_risk,
                  'yearly_expense'    : yearly_expense,
                  'starting_age'      : starting_age,
                  'state_abbrev'      : state_abbrev,
                  'demographic_group' : demographic_group,
                }

    base_save = how_much_to_save(**base_opts)/1e6

    for i, factor in enumerate(factors.keys()):

        opts = base_opts.copy()

        factor_res = []

        for factor_value in factors[factor]['values']:

            opts[factor] = factor_value

            factor_res.append( how_much_to_save(**opts)/1e6 )

        axs[i].plot(factors[factor]['values'], factor_res,
                    marker='.', markersize=3.5,)

        axs[i].plot(base_opts[factor], base_save,
                    marker='+', markersize=5.5,)

        axs[i].set_xlabel(factor)

        #axs[i].legend(fontsize='x-small')
        #plt.ylim(ymin=0, ymax=100)

    axs[1].set_ylabel('Amount to save (million USD)')
    fig.tight_layout()
    fig.savefig('figs/{}.pdf'.format('sensitivity-plots'))

    return fig

if __name__ == '__main__':
    starting_age = 65.0
    state_abbrev = 'IA'
    demographic_group = 'wf'

    # Expenses per year
    yearly_expenses = 50e3

    # Assets
    starting_assets = 3e6

    # Investment allocation
    stock_fraction = 0.5

    # Run one simulation for a given starting assets and stock fraction
    run_histories(starting_assets, yearly_expenses,
                  stock_fraction,
                  starting_age,
                  state_abbrev,
                  demographic_group,
                  n_mc=1000, plotting=True, verbose=True)

    a = how_much_to_save(0.1, yearly_expenses,
                         stock_fraction,
                         starting_age,
                         state_abbrev,
                         demographic_group)

    a = sensitivity_plots()

    """
    # Run many simulations over a range of starting assets and stock fractions
    cascade_plot(yearly_expenses,
                 stock_fraction,
                 starting_age,
                 state_abbrev,
                 demographic_group)
    """
<|endoftext|>


<commit_before>#!/usr/local/bin/python

from datetime import timedelta, datetime
import os
import requests
import re

#
# This will archive inactive channels. The inactive period is in days as 'DAYS_INACTIVE'
# You can put this in a cron job to run daily to do slack cleanup.
#

SLACK_TOKEN      = os.environ.get('SLACK_TOKEN')
DAYS_INACTIVE    = 60
TOO_OLD_DATETIME = datetime.now() - timedelta(days=DAYS_INACTIVE)
WHITELIST_CHANNELS = [r'acct.*', r'.*desk', r'.*engagement.*', r'alert.*']
DRY_RUN = os.environ.get('DRY_RUN')

# api_endpoint is a string, and payload is a dict
def slack_api_http_get(api_endpoint=None, payload=None):
  uri = 'https://slack.com/api/' + api_endpoint
  payload['token'] = SLACK_TOKEN
  try:
    response = requests.get(uri, params=payload)
    if response.status_code == requests.codes.ok:
      return response.json()
    else:
      raise Exception(response.content)
  except Exception as e:
    raise Exception(e)

# too_old_datetime is a datetime object
def get_all_channels():
  payload  = {'exclude_archived': 1}
  api_endpoint = 'channels.list'
  channels = slack_api_http_get(api_endpoint=api_endpoint, payload=payload)['channels']
  all_channels = []
  for channel in channels:
    all_channels.append({'id': channel['id'], 'name': channel['name'], 'created': channel['created'], 'members': channel['members']})
  return all_channels

def get_last_message_timestamp(channel_history, too_old_datetime):
  last_message_datetime = too_old_datetime
  for message in channel_history['messages']:
    if 'subtype' not in message or message['subtype'] == 'file_share' or message['subtype'] == 'file_comment':
      last_message_datetime = datetime.fromtimestamp(float(message['ts']))
      break
  return last_message_datetime

def get_inactive_channels(all_unarchived_channels, too_old_datetime):
  payload  = {'inclusive': 0, 'oldest': 0, 'count': 50}
  api_endpoint = 'channels.history'
  inactive_channels = []
  for channel in all_unarchived_channels:
    payload['channel'] = channel['id']
    channel_history = slack_api_http_get(api_endpoint=api_endpoint, payload=payload)
    last_message_datetime = get_last_message_timestamp(channel_history, datetime.fromtimestamp(float(channel['created'])))
    if last_message_datetime <= too_old_datetime:
      if not (len(channel_history['messages']) > 30 and len(channel['members']) > 5):
        inactive_channels.append(channel)
  return inactive_channels

def filter_out_whitelist_channels(inactive_channels):
    channels_to_archive = []
    for channel in inactive_channels:
      whitelisted = False
      for p in WHITELIST_CHANNELS:
        if re.match(p, channel['name']):
          whitelisted = True
      if not whitelisted:
        channels_to_archive.append(channel)
    return channels_to_archive

def send_channel_message(channel_id, message):
  payload  = {'channel': channel_id, 'username': 'channel_reaper', 'icon_emoji': ':ghost:', 'text': message}
  api_endpoint = 'chat.postMessage'
  slack_api_http_get(api_endpoint=api_endpoint, payload=payload)

def archive_inactive_channels(channels):
  api_endpoint = 'channels.archive'
  for channel in channels:
    if not DRY_RUN:
      message = "This channel has had no activity for %s days. It is being auto-archived." % DAYS_INACTIVE
      message += " If you feel this is a mistake you an unarchive this channel to bring it back at any point."
      send_channel_message(channel['id'], message)
      payload = {'channel': channel['id']}
      slack_api_http_get(api_endpoint=api_endpoint, payload=payload)
    print "Archiving channel... %s" % channel['name']

all_unarchived_channels = get_all_channels()
inactive_channels       = get_inactive_channels(all_unarchived_channels, TOO_OLD_DATETIME)
channels_to_archive = filter_out_whitelist_channels(inactive_channels)
archive_inactive_channels(channels_to_archive)

<commit_msg>Add progress msg<commit_after>#!/usr/local/bin/python

from datetime import timedelta, datetime
import os
import requests
import re

#
# This will archive inactive channels. The inactive period is in days as 'DAYS_INACTIVE'
# You can put this in a cron job to run daily to do slack cleanup.
#

SLACK_TOKEN      = os.environ.get('SLACK_TOKEN')
DAYS_INACTIVE    = 60
TOO_OLD_DATETIME = datetime.now() - timedelta(days=DAYS_INACTIVE)
WHITELIST_CHANNELS = [r'acct.*', r'.*desk', r'.*engagement.*', r'alert.*']
DRY_RUN = os.environ.get('DRY_RUN')

# api_endpoint is a string, and payload is a dict
def slack_api_http_get(api_endpoint=None, payload=None):
  uri = 'https://slack.com/api/' + api_endpoint
  payload['token'] = SLACK_TOKEN
  try:
    response = requests.get(uri, params=payload)
    if response.status_code == requests.codes.ok:
      return response.json()
    else:
      raise Exception(response.content)
  except Exception as e:
    raise Exception(e)

# too_old_datetime is a datetime object
def get_all_channels():
  payload  = {'exclude_archived': 1}
  api_endpoint = 'channels.list'
  channels = slack_api_http_get(api_endpoint=api_endpoint, payload=payload)['channels']
  all_channels = []
  for channel in channels:
    all_channels.append({'id': channel['id'], 'name': channel['name'], 'created': channel['created'], 'members': channel['members']})
  return all_channels

def get_last_message_timestamp(channel_history, too_old_datetime):
  last_message_datetime = too_old_datetime
  for message in channel_history['messages']:
    if 'subtype' not in message or message['subtype'] == 'file_share' or message['subtype'] == 'file_comment':
      last_message_datetime = datetime.fromtimestamp(float(message['ts']))
      break
  return last_message_datetime

def get_inactive_channels(all_unarchived_channels, too_old_datetime):
  print "Find inactive channels..."
  payload  = {'inclusive': 0, 'oldest': 0, 'count': 50}
  api_endpoint = 'channels.history'
  inactive_channels = []
  for channel in all_unarchived_channels:
    payload['channel'] = channel['id']
    channel_history = slack_api_http_get(api_endpoint=api_endpoint, payload=payload)
    last_message_datetime = get_last_message_timestamp(channel_history, datetime.fromtimestamp(float(channel['created'])))
    if last_message_datetime <= too_old_datetime:
      if not (len(channel_history['messages']) > 30 and len(channel['members']) > 5):
        inactive_channels.append(channel)
  return inactive_channels

def filter_out_whitelist_channels(inactive_channels):
    channels_to_archive = []
    for channel in inactive_channels:
      whitelisted = False
      for p in WHITELIST_CHANNELS:
        if re.match(p, channel['name']):
          whitelisted = True
      if not whitelisted:
        channels_to_archive.append(channel)
    return channels_to_archive

def send_channel_message(channel_id, message):
  payload  = {'channel': channel_id, 'username': 'channel_reaper', 'icon_emoji': ':ghost:', 'text': message}
  api_endpoint = 'chat.postMessage'
  slack_api_http_get(api_endpoint=api_endpoint, payload=payload)

def archive_inactive_channels(channels):
  print "Archive inactive channels..."
  api_endpoint = 'channels.archive'
  for channel in channels:
    if not DRY_RUN:
      message = "This channel has had no activity for %s days. It is being auto-archived." % DAYS_INACTIVE
      message += " If you feel this is a mistake you an unarchive this channel to bring it back at any point."
      send_channel_message(channel['id'], message)
      payload = {'channel': channel['id']}
      slack_api_http_get(api_endpoint=api_endpoint, payload=payload)
    print "Archiving channel... %s" % channel['name']

all_unarchived_channels = get_all_channels()
inactive_channels       = get_inactive_channels(all_unarchived_channels, TOO_OLD_DATETIME)
channels_to_archive = filter_out_whitelist_channels(inactive_channels)
archive_inactive_channels(channels_to_archive)

<|endoftext|>


<commit_before>from disco.test import TestCase, TestJob
from disco.core import Job
import disco
import threading
import BaseHTTPServer

def map(line, params):
    for word in line.split():
        yield word, 1

def reduce(iter, params):
    from disco.util import kvgroup
    for word, counts in kvgroup(sorted(iter)):
        yield word, sum(counts)

PORT = 1234

class MyHandler(BaseHTTPServer.BaseHTTPRequestHandler):
    def do_GET(s):
        s.send_response(200)
        s.send_header("Content-type", "text/html")
        s.send_header("Transfer-Encoding", "chunked")
        s.end_headers()
        s.wfile.write("b\r\nHello World\r\n0\r\n\r\n")

def startServer():
    server_class = BaseHTTPServer.HTTPServer
    httpd = server_class(('', PORT), MyHandler)
    httpd.handle_request()
    httpd.server_close()

class RawTestCase(TestCase):
    def runTest(self):
        threading.Thread(target=startServer).start()
        input = 'http:' + self.disco.master.split(':')[1] + ":" + str(PORT)
        self.job = Job().run(input=[input], map=map, reduce=reduce)
        self.assertEqual(sorted(self.results(self.job)), [('Hello', 1), ('World', 1)])
<commit_msg>Use the disco.compat.http_server to work with python3.<commit_after>from disco.test import TestCase, TestJob
from disco.core import Job
from disco.compat import http_server
import disco
import threading

def map(line, params):
    for word in line.split():
        yield word, 1

def reduce(iter, params):
    from disco.util import kvgroup
    for word, counts in kvgroup(sorted(iter)):
        yield word, sum(counts)

PORT = 1234

class MyHandler(http_server.BaseHTTPRequestHandler):
    def do_GET(s):
        s.send_response(200)
        s.send_header("Content-type", "text/html")
        s.send_header("Transfer-Encoding", "chunked")
        s.end_headers()
        s.wfile.write("b\r\nHello World\r\n0\r\n\r\n")

def startServer():
    server_class = http_server.HTTPServer
    httpd = server_class(('', PORT), MyHandler)
    httpd.handle_request()
    httpd.server_close()

class RawTestCase(TestCase):
    def runTest(self):
        threading.Thread(target=startServer).start()
        input = 'http:' + self.disco.master.split(':')[1] + ":" + str(PORT)
        self.job = Job().run(input=[input], map=map, reduce=reduce)
        self.assertEqual(sorted(self.results(self.job)), [('Hello', 1), ('World', 1)])
<|endoftext|>


<commit_before><commit_msg>extend task duration<commit_after><|endoftext|>


<commit_before>"""Queue worker and manager

Any job that takes keyword arguments can be a ``job`` that a worker runs. The
:class:`~retools.queue.QueueManager` handles configuration and enqueing jobs
to be run.

Declaring jobs::

    # jobs.py
    def default_job():
        # do some basic thing

    def important(somearg=None):
        # do an important thing

    def my_event_handler(sender, **kwargs):
        # do something

    def save_error(sender, **kwargs):
        # record error

    def configure_events(event):
        # Hook up the my_event_handler just for the default_job function
        event.listen('job_postrun', my_event_handler, jobs=[default_job])

        # Hook up a handler for all job failures
        event.listen('job_failure', save_error)

Running Jobs::

    # main.py
    from retools.queue import QueueManager

    qm = QueueManager()
    qm.enqueue('mypackage.jobs:important', somearg='fred')

The :meth:`~retools.queue.QueueManager.scan` must be called before jobs can
be enqueued.

"""
import os
import signal
import socket
import subprocess
import sys
import time
import uuid
from datetime import datetime
from optparse import OptionParser

import pkg_resources

try:
    import json
except ImportError: #pragma: nocover
    import simplejson as json

from setproctitle import setproctitle

from retools import global_connection
from retools.event import Event
from retools.exc import UnregisteredJob
from retools.exc import ConfigurationError
from retools.util import with_nested_contexts

class QueueManager(object):
    """Configures and enqueues jobs"""
    def __init__(self, redis=None, default_queue_name='main'):
        """Initialize a QueueManager

        :param redis: A Redis instance. Defaults to the redis instance
                      on the global_connection.
        """
        self.default_queue_name = default_queue_name
        self.redis = redis or global_connection.redis
        self.names = {} # cache name lookups
        self.job_config = {}

    def set_queue_for_job(self, job_name, queue_name):
        """Set the queue that a given job name will go to

        :param job_name: The pkg_resource name of the job function. I.e.
                         retools.jobs:my_function
        :param queue_name: Name of the queue on Redis job payloads should go
                           to
        """
        self.job_config[job_name] = queue_name

    def enqueue(self, job, **kwargs):
        """Enqueue a job

        :param job: The pkg_resouce name of the function. I.e.
                    retools.jobs:my_function
        :param kwargs: Keyword arguments the job should be called with.
                       These arguments must be serializeable by JSON.
        :returns: The job id that was queued.

        """
        if job not in self.names:
            job_func = pkg_resources.EntryPoint.parse('x=%s' % job).load(False)
            self.names[job] = job_func

        queue_name = kwargs.pop('queue_name', None)
        if not queue_name:
            queue_name = self.job_config.get('job', self.default_queue_name)

        full_queue_name = 'retools:queue:' + queue_name
        job_id = uuid.uuid4().hex
        job_dct = {
            'job_id': job_id,
            'job': job,
            'kwargs': kwargs,
            'state': {}
        }
        pipeline = self.redis.pipeline()
        pipeline.rpush(full_queue_name, json.dumps(job_dct))
        pipeline.sadd('retools:queues', queue_name)
        pipeline.execute()
        return job_id

class Job(object):
    def __init__(self, queue_name, job_payload, redis, event):
        """Create a job instance given a JSON job payload

        :param job_payload: A JSON string representing a job.
        :param queue_name: The queue this job was pulled off of.
        :param redis: The redis instance used to pull this job.
        :param event: An :class:`~retools.event.Event` instance used to
                      dispatch events.

        A ``Job`` instance is created when the Worker pulls a
        job payload off the queue. The ``current_job`` global is set
        upon creation to indicate the current job being processed.

        """
        global current_job
        current_job = self

        self.event = event
        self.payload = payload = json.loads(job_payload)
        self.job_id = payload['job_id']
        self.job_name = payload['job']
        self.queue_name = queue_name
        self.kwargs = payload['kwargs']
        self.state = payload['state']
        self.redis = redis
        self.func = None

    def perform(self):
        """Runs the job calling all the job signals as appropriate"""
        event = self.event
        event['job_prerun'].send(self.func, job=self)
        try:
            if event['job_wrapper'].has_receivers_for(self.func):
                handlers = list(event['job_wrapper'].recievers_for(self.func))
                result = with_nested_contexts(handlers, self.func, [self], self.kwargs)
            else:
                result = self.func(**self.kwargs)
            event['job_postrun'].send(self.func, job=self, result=result)
            return True
        except Exception, exc:
            event['job_failure'].send(self.func, job=self, exc=exc)
            return False

    def enqueue(self):
        full_queue_name = self.queue_name
        job_dct = {
            'job_id': self.job_id,
            'job': self.job_name,
            'kwargs': self.kwargs,
            'state': self.state
        }
        queue_name = full_queue_name.lstrip('retools:queue:')
        pipeline = self.redis.pipeline()
        pipeline.rpush(full_queue_name, json.dumps(job_dct))
        pipeline.sadd('retools:queues', queue_name)
        pipeline.execute()
        return job_id

class Worker(object):
    """A Worker works on jobs"""
    def __init__(self, queues, event, redis=None):
        """Create a worker

        :param queues: List of queues to process
        :type queues: list
        :param event: An :class:`~retools.event.Event` instance used to
                      dispatch events.
        :param redis: Redis instance to use, defaults to the global_connection.

        In the event that there is only a single queue in the list
        Redis list blocking will be used for lower latency job
        processing

        """
        self.event = event
        self.redis = redis or global_connection.redis
        if not queues:
            raise ConfigurationError("No queues were configured for this worker")
        self.queues = ['retools:queue:%s' % x for x in queues]
        self.paused = self.shutdown = False
        self.job = None
        self.child_id = None
        self.jobs = {} # job function import cache

    @property
    def worker_id(self):
        """Returns this workers id based on hostname, pid, queues"""
        return '%s:%s:%s' % (socket.gethostname(), os.getpid(), self.queue_names)

    @property
    def queue_names(self):
        names = [x.lstrip('retools:queue:') for x in self.queues]
        return ','.join(names)

    def work(self, interval=5, blocking=False):
        """Work on jobs

        This is the main method of the Worker, and will register itself
        with Redis as a Worker, wait for jobs, then process them.

        :param interval: Time in seconds between polling.
        :type interval: int
        :param blocking: Whether or not blocking pop should be used. If the
                         blocking pop is used, then the worker will block for
                         ``interval`` seconds at a time waiting for a new
                         job. This affects how often the worke can respond to
                         OS signals.
        :type blocking: bool

        """
        self.set_proc_title('Starting')
        self.startup()

        try:
            while 1:
                if self.shutdown:
                    break

                # Set this first since reserve may block for awhile
                self.set_proc_title("Waiting for %s" % self.queue_names)

                if not self.paused and self.reserve(interval, blocking):
                    self.event['worker_prefork'].send(self, job=self.job)
                    self.working_on()
                    self.child_id = os.fork()
                    if self.child_id:
                        self.set_proc_title("Forked %s at %s" % (
                            self.child_id, datetime.now()))
                        os.wait()
                    else:
                        self.set_proc_title("Processing %s since %s" % (
                            self.job.queue_name, datetime.now()))
                        self.perform()
                        sys.exit()
                    self.done_working()
                    self.child_id = None
                    self.job = None
                else:
                    if self.paused:
                        self.set_proc_title("Paused")
                    elif not blocking:
                        self.set_proc_title("Waiting for %s" % self.queue_names)
                        time.sleep(interval)
        finally:
            self.unregister_worker()
            self.event['worker_shutdown'].send(self)

    def reserve(self, interval, blocking):
        """Attempts to pull a job off the queue(s)"""
        queue_name = None
        if blocking:
            result = self.redis.blpop(self.queues, timeout=interval)
            if result:
                queue_name, job_payload = result
        else:
            for queue in self.queues:
                job_payload = self.redis.lpop(queue)
                if job_payload:
                    queue_name = queue
                    break
        if not queue_name:
            return False

        self.job = job = Job(queue_name=queue_name, job_payload=job_payload,
                             redis=self.redis, event=self.event)
        try:
            job.func = self.jobs[job.job_name]
        except KeyError:
            mod_name, func_name = job.job_name.split(':')
            __import__(mod_name)
            mod = sys.modules[mod_name]
            configure_events = getattr(mod, 'configure_events', None)
            if configure_events:
                configure_events(self.event)
            job.func = self.jobs[job.job_name] = getattr(mod, func_name)
        return True

    def set_proc_title(self, title):
        """Sets the active process title, retains the retools prefic"""
        setproctitle('retools: ' + title)

    def register_worker(self):
        """Register this worker with Redis"""
        pipeline = self.redis.pipeline()
        pipeline.sadd("retools:workers", self.worker_id)
        pipeline.set("retools:worker:%s:started" % self.worker_id, time.time())
        pipeline.execute()

    def unregister_worker(self, worker_id=None):
        """Unregister this worker with Redis"""
        worker_id = worker_id or self.worker_id
        pipeline = self.redis.pipeline()
        pipeline.srem("retools:workers", worker_id)
        pipeline.delete("retools:worker:%s" % worker_id)
        pipeline.delete("retools:worker:%s:started" % worker_id)
        pipeline.execute()

    def startup(self):
        """Runs basic startup tasks"""
        self.register_signal_handlers()
        self.prune_dead_workers()
        self.register_worker()
        self.event['worker_startup'].send(self)

    def trigger_shutdown(self, *args):
        """Graceful shutdown of the worker"""
        self.shutdown = True

    def immediate_shutdown(self, *args):
        """Immediately shutdown the worker, kill child process if needed"""
        self.shutdown = True
        self.kill_child()

    def kill_child(self, *args):
        """Kill the child process immediately"""
        if self.child_id:
            os.kill(self.child_id, signal.SIGTERM)

    def pause_processing(self, *args):
        """Cease pulling jobs off the queue for processing"""
        self.paused = True

    def resume_processing(self, *args):
        """Resume pulling jobs for processing off the queue"""
        self.paused = False

    def prune_dead_workers(self):
        """Prune dead workers from Redis"""
        all_workers = self.redis.smembers("retools:workers")
        known_workers = self.worker_pids()
        hostname = socket.gethostname()
        for worker in all_workers:
            host, pid, queues = worker.split(':')
            if host != hostname or pid in known_workers:
                continue
            self.unregister_worker(worker)

    def register_signal_handlers(self):
        """Setup all the signal handlers"""
        signal.signal(signal.SIGTERM, self.immediate_shutdown)
        signal.signal(signal.SIGINT, self.immediate_shutdown)
        signal.signal(signal.SIGQUIT, self.trigger_shutdown)
        signal.signal(signal.SIGUSR1, self.kill_child)
        signal.signal(signal.SIGUSR2, self.pause_processing)
        signal.signal(signal.SIGCONT, self.resume_processing)

    def working_on(self):
        """Indicate with Redis what we're working on"""
        data = {
            'queue': self.job.queue_name,
            'run_at': time.time(),
            'payload': self.job.payload
        }
        self.redis.set("retools:worker:%s" % self.worker_id, json.dumps(data))

    def done_working(self):
        """Called when we're done working on a job"""
        self.redis.delete("retools:worker:%s" % self.worker_id)

    def worker_pids(self):
        """Returns a list of all the worker processes"""
        ps = subprocess.Popen("ps -U 0 -A | grep 'retools:'", shell=True,
                              stdout=subprocess.PIPE)
        data = ps.stdout.read()
        ps.stdout.close()
        ps.wait()
        return [x.split()[0] for x in data.split('\n') if x]

    def perform(self):
        """Run the job and call the appropriate signal handlers"""
        self.event['worker_postfork'].send(self, job=self.job)
        self.job.perform()

def run_worker():
    usage = "usage: %prog queues"
    parser = OptionParser(usage=usage)
    parser.add_option("--interval", dest="interval", type="int", default=5,
                      help="Polling interval")
    parser.add_option("-b", dest="blocking", action="store_true", default=False,
                      help="Whether to use blocking queue semantics")
    (options, args) = parser.parse_args()

    if len(args) < 1:
        sys.exit("Error: Failed to provide queues or packages_to_scan args")

    event = Event()
    worker = Worker(queues=args[0].split(','), event=event)
    worker.work(interval=options.interval, blocking=options.blocking)
    sys.exit()
<commit_msg>Rework events as they will now be tracked with the job up-front<commit_after>"""Queue worker and manager

Any job that takes keyword arguments can be a ``job`` that a worker runs. The
:class:`~retools.queue.QueueManager` handles configuration and enqueing jobs
to be run.

Declaring jobs::

    # jobs.py
    def default_job():
        # do some basic thing

    def important(somearg=None):
        # do an important thing

    def my_event_handler(sender, **kwargs):
        # do something

    def save_error(sender, **kwargs):
        # record error

Running Jobs::

    # main.py
    from retools.queue import QueueManager

    qm = QueueManager()
    qm.subscriber('job_failure', handler='mypackage.jobs:save_error')
    qm.subscriber('job_postrun', 'mypackage.jobs:important',
                  handler='mypackage.jobs:my_event_handler)
    qm.enqueue('mypackage.jobs:important', somearg='fred')

The :meth:`~retools.queue.QueueManager.scan` must be called before jobs can
be enqueued.

"""
import os
import signal
import socket
import subprocess
import sys
import time
import uuid
from datetime import datetime
from optparse import OptionParser

import pkg_resources

try:
    import json
except ImportError: #pragma: nocover
    import simplejson as json

from setproctitle import setproctitle

from retools import global_connection
from retools.event import Event
from retools.exc import UnregisteredJob
from retools.exc import ConfigurationError
from retools.util import with_nested_contexts

class QueueManager(object):
    """Configures and enqueues jobs"""
    def __init__(self, redis=None, default_queue_name='main'):
        """Initialize a QueueManager

        :param redis: A Redis instance. Defaults to the redis instance
                      on the global_connection.
        """
        self.default_queue_name = default_queue_name
        self.redis = redis or global_connection.redis
        self.names = {} # cache name lookups
        self.job_config = {}
        self.job_events = {}
        self.global_events = {}

    def set_queue_for_job(self, job_name, queue_name):
        """Set the queue that a given job name will go to

        :param job_name: The pkg_resource name of the job function. I.e.
                         retools.jobs:my_function
        :param queue_name: Name of the queue on Redis job payloads should go
                           to
        """
        self.job_config[job_name] = queue_name

    def subscriber(self, event, job=None, handler=None):
        """Set events for a specific job or for all jobs

        :param event: The name of the event to subscribe to.
        :param job: Optional, a specific job to bind to.
        :param handler: The location of the handler to call.

        """
        if job:
            self.job_events.setdefault(event, []).append(handler)
        else:
            self.global_events.setdefault(event, []).append(handler)

    def enqueue(self, job, **kwargs):
        """Enqueue a job

        :param job: The pkg_resouce name of the function. I.e.
                    retools.jobs:my_function
        :param kwargs: Keyword arguments the job should be called with.
                       These arguments must be serializeable by JSON.
        :returns: The job id that was queued.

        """
        if job not in self.names:
            job_func = pkg_resources.EntryPoint.parse('x=%s' % job).load(False)
            self.names[job] = job_func

        queue_name = kwargs.pop('queue_name', None)
        if not queue_name:
            queue_name = self.job_config.get('job', self.default_queue_name)

        full_queue_name = 'retools:queue:' + queue_name
        job_id = uuid.uuid4().hex
        events = self.global_events.copy()
        if job in self.job_events:
            for k, v in self.job_events[job].items():
                if k in events:
                    events[k].extend(v)
                else:
                    events[k] = v

        job_dct = {
            'job_id': job_id,
            'job': job,
            'kwargs': kwargs,
            'events': events,
            'state': {}
        }
        pipeline = self.redis.pipeline()
        pipeline.rpush(full_queue_name, json.dumps(job_dct))
        pipeline.sadd('retools:queues', queue_name)
        pipeline.execute()
        return job_id

class Job(object):
    def __init__(self, queue_name, job_payload, redis, event):
        """Create a job instance given a JSON job payload

        :param job_payload: A JSON string representing a job.
        :param queue_name: The queue this job was pulled off of.
        :param redis: The redis instance used to pull this job.
        :param event: An :class:`~retools.event.Event` instance used to
                      dispatch events.

        A ``Job`` instance is created when the Worker pulls a
        job payload off the queue. The ``current_job`` global is set
        upon creation to indicate the current job being processed.

        """
        global current_job
        current_job = self

        self.event = event
        self.payload = payload = json.loads(job_payload)
        self.job_id = payload['job_id']
        self.job_name = payload['job']
        self.queue_name = queue_name
        self.kwargs = payload['kwargs']
        self.state = payload['state']
        self.events = {}
        self.redis = redis
        self.func = None
        self.load_events()

    def load_events(self):
        """Load all the events given the references"""
        for k, v in self.payload['events'].items():
            funcs = []
            for name in v:
                mod_name, func_name = name.split(':')
                try:
                    mod = sys.modules[mod_name]
                except KeyError:
                    __import__(mod_name)
                    mod = sys.modules[mod_name]
                funcs.append(getattr(mod, func_name))
            self.events[k] = funcs

    def perform(self):
        """Runs the job calling all the job signals as appropriate"""
        event = self.event
        event['job_prerun'].send(self.func, job=self)
        try:
            if event['job_wrapper'].has_receivers_for(self.func):
                handlers = list(event['job_wrapper'].recievers_for(self.func))
                result = with_nested_contexts(handlers, self.func, [self], self.kwargs)
            else:
                result = self.func(**self.kwargs)
            event['job_postrun'].send(self.func, job=self, result=result)
            return True
        except Exception, exc:
            event['job_failure'].send(self.func, job=self, exc=exc)
            return False

    def enqueue(self):
        full_queue_name = self.queue_name
        job_dct = {
            'job_id': self.job_id,
            'job': self.job_name,
            'kwargs': self.kwargs,
            'events': self.payload['events'],
            'state': self.state
        }
        queue_name = full_queue_name.lstrip('retools:queue:')
        pipeline = self.redis.pipeline()
        pipeline.rpush(full_queue_name, json.dumps(job_dct))
        pipeline.sadd('retools:queues', queue_name)
        pipeline.execute()
        return job_id

    def run_event(self, event):
        pass

class Worker(object):
    """A Worker works on jobs"""
    def __init__(self, queues, event, redis=None):
        """Create a worker

        :param queues: List of queues to process
        :type queues: list
        :param event: An :class:`~retools.event.Event` instance used to
                      dispatch events.
        :param redis: Redis instance to use, defaults to the global_connection.

        In the event that there is only a single queue in the list
        Redis list blocking will be used for lower latency job
        processing

        """
        self.event = event
        self.redis = redis or global_connection.redis
        if not queues:
            raise ConfigurationError("No queues were configured for this worker")
        self.queues = ['retools:queue:%s' % x for x in queues]
        self.paused = self.shutdown = False
        self.job = None
        self.child_id = None
        self.jobs = {} # job function import cache

    @property
    def worker_id(self):
        """Returns this workers id based on hostname, pid, queues"""
        return '%s:%s:%s' % (socket.gethostname(), os.getpid(), self.queue_names)

    @property
    def queue_names(self):
        names = [x.lstrip('retools:queue:') for x in self.queues]
        return ','.join(names)

    def work(self, interval=5, blocking=False):
        """Work on jobs

        This is the main method of the Worker, and will register itself
        with Redis as a Worker, wait for jobs, then process them.

        :param interval: Time in seconds between polling.
        :type interval: int
        :param blocking: Whether or not blocking pop should be used. If the
                         blocking pop is used, then the worker will block for
                         ``interval`` seconds at a time waiting for a new
                         job. This affects how often the worke can respond to
                         OS signals.
        :type blocking: bool

        """
        self.set_proc_title('Starting')
        self.startup()

        try:
            while 1:
                if self.shutdown:
                    break

                # Set this first since reserve may block for awhile
                self.set_proc_title("Waiting for %s" % self.queue_names)

                if not self.paused and self.reserve(interval, blocking):
                    self.working_on()
                    self.child_id = os.fork()
                    if self.child_id:
                        self.set_proc_title("Forked %s at %s" % (
                            self.child_id, datetime.now()))
                        os.wait()
                    else:
                        self.set_proc_title("Processing %s since %s" % (
                            self.job.queue_name, datetime.now()))
                        self.perform()
                        sys.exit()
                    self.done_working()
                    self.child_id = None
                    self.job = None
                else:
                    if self.paused:
                        self.set_proc_title("Paused")
                    elif not blocking:
                        self.set_proc_title("Waiting for %s" % self.queue_names)
                        time.sleep(interval)
        finally:
            self.unregister_worker()

    def reserve(self, interval, blocking):
        """Attempts to pull a job off the queue(s)"""
        queue_name = None
        if blocking:
            result = self.redis.blpop(self.queues, timeout=interval)
            if result:
                queue_name, job_payload = result
        else:
            for queue in self.queues:
                job_payload = self.redis.lpop(queue)
                if job_payload:
                    queue_name = queue
                    break
        if not queue_name:
            return False

        self.job = job = Job(queue_name=queue_name, job_payload=job_payload,
                             redis=self.redis, event=self.event)
        try:
            job.func = self.jobs[job.job_name]
        except KeyError:
            mod_name, func_name = job.job_name.split(':')
            __import__(mod_name)
            mod = sys.modules[mod_name]
            configure_events = getattr(mod, 'configure_events', None)
            if configure_events:
                configure_events(self.event)
            job.func = self.jobs[job.job_name] = getattr(mod, func_name)
        return True

    def set_proc_title(self, title):
        """Sets the active process title, retains the retools prefic"""
        setproctitle('retools: ' + title)

    def register_worker(self):
        """Register this worker with Redis"""
        pipeline = self.redis.pipeline()
        pipeline.sadd("retools:workers", self.worker_id)
        pipeline.set("retools:worker:%s:started" % self.worker_id, time.time())
        pipeline.execute()

    def unregister_worker(self, worker_id=None):
        """Unregister this worker with Redis"""
        worker_id = worker_id or self.worker_id
        pipeline = self.redis.pipeline()
        pipeline.srem("retools:workers", worker_id)
        pipeline.delete("retools:worker:%s" % worker_id)
        pipeline.delete("retools:worker:%s:started" % worker_id)
        pipeline.execute()

    def startup(self):
        """Runs basic startup tasks"""
        self.register_signal_handlers()
        self.prune_dead_workers()
        self.register_worker()

    def trigger_shutdown(self, *args):
        """Graceful shutdown of the worker"""
        self.shutdown = True

    def immediate_shutdown(self, *args):
        """Immediately shutdown the worker, kill child process if needed"""
        self.shutdown = True
        self.kill_child()

    def kill_child(self, *args):
        """Kill the child process immediately"""
        if self.child_id:
            os.kill(self.child_id, signal.SIGTERM)

    def pause_processing(self, *args):
        """Cease pulling jobs off the queue for processing"""
        self.paused = True

    def resume_processing(self, *args):
        """Resume pulling jobs for processing off the queue"""
        self.paused = False

    def prune_dead_workers(self):
        """Prune dead workers from Redis"""
        all_workers = self.redis.smembers("retools:workers")
        known_workers = self.worker_pids()
        hostname = socket.gethostname()
        for worker in all_workers:
            host, pid, queues = worker.split(':')
            if host != hostname or pid in known_workers:
                continue
            self.unregister_worker(worker)

    def register_signal_handlers(self):
        """Setup all the signal handlers"""
        signal.signal(signal.SIGTERM, self.immediate_shutdown)
        signal.signal(signal.SIGINT, self.immediate_shutdown)
        signal.signal(signal.SIGQUIT, self.trigger_shutdown)
        signal.signal(signal.SIGUSR1, self.kill_child)
        signal.signal(signal.SIGUSR2, self.pause_processing)
        signal.signal(signal.SIGCONT, self.resume_processing)

    def working_on(self):
        """Indicate with Redis what we're working on"""
        data = {
            'queue': self.job.queue_name,
            'run_at': time.time(),
            'payload': self.job.payload
        }
        self.redis.set("retools:worker:%s" % self.worker_id, json.dumps(data))

    def done_working(self):
        """Called when we're done working on a job"""
        self.redis.delete("retools:worker:%s" % self.worker_id)

    def worker_pids(self):
        """Returns a list of all the worker processes"""
        ps = subprocess.Popen("ps -U 0 -A | grep 'retools:'", shell=True,
                              stdout=subprocess.PIPE)
        data = ps.stdout.read()
        ps.stdout.close()
        ps.wait()
        return [x.split()[0] for x in data.split('\n') if x]

    def perform(self):
        """Run the job and call the appropriate signal handlers"""
        self.event['worker_postfork'].send(self, job=self.job)
        self.job.perform()

def run_worker():
    usage = "usage: %prog queues"
    parser = OptionParser(usage=usage)
    parser.add_option("--interval", dest="interval", type="int", default=5,
                      help="Polling interval")
    parser.add_option("-b", dest="blocking", action="store_true", default=False,
                      help="Whether to use blocking queue semantics")
    (options, args) = parser.parse_args()

    if len(args) < 1:
        sys.exit("Error: Failed to provide queues or packages_to_scan args")

    event = Event()
    worker = Worker(queues=args[0].split(','), event=event)
    worker.work(interval=options.interval, blocking=options.blocking)
    sys.exit()
<|endoftext|>


<commit_before>from cassandra.cluster import Cluster
from pandas import DataFrame, read_fwf
import pytest
from ironeagle import get_prepared_statement, save_dataframe_to_cassandra

@pytest.fixture
def session():
    cluster = Cluster()
    session = cluster.connect("ironeagle")

    create = """
              CREATE TABLE IF NOT EXISTS simple (
              id int primary key,
              name text,
              lat float
              )
              """

    truncate = "truncate simple"

    create_station = """CREATE TABLE if not exists station (
                    station_id text primary key,
                    lat float,
                    long float,
                    elevation float,
                    state text,
                    name text,
                    gsn_flag text,
                    hcn_flag text,
                    wmo_id float
                )
                """

    trunate_station = "truncate station"

    queries = [create, truncate, create_station, trunate_station]

    for q in queries:
        session.execute(q)
    return session

@pytest.fixture
def df():
    df = DataFrame.from_records([(1, "jon", 1.1), (2, "pete", 1.2)],
                                    columns=['id', 'name', 'lat'])
    return df

@pytest.fixture
def stations():
    stations = read_fwf("tests/ghcnd-stations.txt", header=None,
                        colspecs=[(0,11), (12, 20), (21, 30), (31,37), (38,40), (41,71), (72,75), (76,79), (80,85)],
                        names=["station_id", "lat", "long", "elevation", "state", "name", "gsn_flag", "hcn_flag", "wmo_id"])
    return stations

def test_get_prepared_statement(df):
    prepared = get_prepared_statement(df, "simple")
    assert "id, name" in prepared
    assert "(?,?,?)" in prepared

def test_insert_data(df, session):
    save_dataframe_to_cassandra(session, df, "simple")

def test_station_data(session, stations):
    save_dataframe_to_cassandra(session, stations, "station")
<commit_msg>using type conversion to make sure i conform to my type spec<commit_after>from cassandra.cluster import Cluster
from pandas import DataFrame, read_fwf
import pytest
from ironeagle import get_prepared_statement, save_dataframe_to_cassandra

@pytest.fixture
def session():
    cluster = Cluster()
    session = cluster.connect("ironeagle")

    create = """
              CREATE TABLE IF NOT EXISTS simple (
              id int primary key,
              name text,
              lat float
              )
              """

    truncate = "truncate simple"

    create_station = """CREATE TABLE if not exists station (
                    station_id text primary key,
                    lat float,
                    long float,
                    elevation float,
                    state text,
                    name text,
                    gsn_flag text,
                    hcn_flag text,
                    wmo_id text
                )
                """

    trunate_station = "truncate station"

    queries = [create, truncate, create_station, trunate_station]

    for q in queries:
        session.execute(q)
    return session

@pytest.fixture
def df():
    df = DataFrame.from_records([(1, "jon", 1.1), (2, "pete", 1.2)],
                                    columns=['id', 'name', 'lat'])
    return df

@pytest.fixture
def stations():
    stations = read_fwf("tests/ghcnd-stations.txt", header=None,
                        colspecs=[(0,11), (12, 20), (21, 30), (31,37), (38,40), (41,71), (72,75), (76,79), (80,85)],
                        names=["station_id", "lat", "long", "elevation", "state", "name", "gsn_flag", "hcn_flag", "wmo_id"])

    stations['wmo_id'] = stations['wmo_id'].astype(str)
    return stations

def test_get_prepared_statement(df):
    prepared = get_prepared_statement(df, "simple")
    assert "id, name" in prepared
    assert "(?,?,?)" in prepared

def test_insert_data(df, session):
    save_dataframe_to_cassandra(session, df, "simple")

def test_station_data(session, stations):
    save_dataframe_to_cassandra(session, stations, "station")
<|endoftext|>


<commit_before><commit_msg>no message<commit_after><|endoftext|>


<commit_before>from .particles_source import ParticlesSource
from .lattice_source import LatticeSource
from .utils import cell_array_slicer

__all__ = [
    'ParticlesSource',
    'LatticeSource',
    'cell_array_slicer']
<commit_msg>minor cleanup<commit_after><|endoftext|>


<commit_before>from django.contrib import admin
from django.contrib import messages

from .models import Review
from .utils import SteamStore, SteamException

class ReviewAdmin(admin.ModelAdmin):
    list_display = (
        'updated_at', 'steam_app_id', 'name', 'localized_by_developer', 'localized_by_community', 'published', 'invalid'
    )
    list_filter = [
        'updated_at', 'localized_by_developer', 'localized_by_community', 'published', 'invalid'
    ]
    actions = ['mark_as_draft', 'mark_as_published', 'sync_with_steam']

    def mark_as_draft(self, request, queryset):
        rows_updated = queryset.update(published=False)
        if rows_updated == 1:
            message_bit = '1 review was'
        else:
            message_bit = '%s reviews were' % rows_updated
        self.message_user(request, '%s successfully marked as draft.' % message_bit)

    def mark_as_published(self, request, queryset):
        rows_updated = queryset.update(published=True)
        if rows_updated == 1:
            message_bit = '1 review was'
        else:
            message_bit = '%s reviews were' % rows_updated
        self.message_user(request, '%s successfully marked as published.' % message_bit)

    def sync_with_steam(self, request, queryset):
        store = SteamStore()
        success = 0

        for row in queryset:
            data = None
            app_id = row.steam_app_id

            try:
                data = store.appdetails(app_id)
            except SteamException as e:
                self.message_user(request, 'Steam AppID {:d} is invalid. ({:s})'.format(app_id, str(e)), messages.ERROR)
                row.invalid = True

            if data:
                row.name = store.name(data)
                row.localized_by_developer = store.is_support_japanese(data)
                row.invalid = False
                success += 1

            row.save()

        if success >= 1:
            if success == 1:
                message_bit = '1 review was'
            else:
                message_bit = '%s reviews were' % success
            self.message_user(request, '%s successfully synced with Steam.' % message_bit)

admin.site.register(Review, ReviewAdmin)
<commit_msg>show all fields in Review on admin site<commit_after>from django.contrib import admin
from django.contrib import messages

from .models import Review
from .utils import SteamStore, SteamException

class ReviewAdmin(admin.ModelAdmin):
    list_display = (
        'updated_at', 'steam_app_id', 'name', 'localization_status', 'review_summary', 'review_detail', 'localized_by_developer', 'localized_by_community', 'published', 'invalid'
    )
    list_filter = [
        'updated_at', 'localized_by_developer', 'localized_by_community', 'published', 'invalid'
    ]
    actions = ['mark_as_draft', 'mark_as_published', 'sync_with_steam']

    def mark_as_draft(self, request, queryset):
        rows_updated = queryset.update(published=False)
        if rows_updated == 1:
            message_bit = '1 review was'
        else:
            message_bit = '%s reviews were' % rows_updated
        self.message_user(request, '%s successfully marked as draft.' % message_bit)

    def mark_as_published(self, request, queryset):
        rows_updated = queryset.update(published=True)
        if rows_updated == 1:
            message_bit = '1 review was'
        else:
            message_bit = '%s reviews were' % rows_updated
        self.message_user(request, '%s successfully marked as published.' % message_bit)

    def sync_with_steam(self, request, queryset):
        store = SteamStore()
        success = 0

        for row in queryset:
            data = None
            app_id = row.steam_app_id

            try:
                data = store.appdetails(app_id)
            except SteamException as e:
                self.message_user(request, 'Steam AppID {:d} is invalid. ({:s})'.format(app_id, str(e)), messages.ERROR)
                row.invalid = True

            if data:
                row.name = store.name(data)
                row.localized_by_developer = store.is_support_japanese(data)
                row.invalid = False
                success += 1

            row.save()

        if success >= 1:
            if success == 1:
                message_bit = '1 review was'
            else:
                message_bit = '%s reviews were' % success
            self.message_user(request, '%s successfully synced with Steam.' % message_bit)

admin.site.register(Review, ReviewAdmin)
<|endoftext|>