<commit_before>import json
import typing
from operator import attrgetter

from django.contrib import messages
from django.contrib.auth import get_user_model
from django.contrib.contenttypes.models import ContentType
from django.core.exceptions import PermissionDenied
from django.http import HttpRequest, HttpResponse, JsonResponse
from django.shortcuts import render, redirect
from django.urls import reverse
from django.views import View

from accounts.db_facade import fetch_team_for_account
from accounts.forms import TeamForm
from accounts.models import Team, AccountPermissionType
from accounts.views import AccountPermissionsMixin
from projects.permission_models import ProjectRole, ProjectAgentRole, AgentType
from projects.project_models import Project

User = get_user_model()

AGENT_ROLE_ID_PREFIX = 'agent_role_id_'

class TeamDetailView(AccountPermissionsMixin, View):
    def get(self, request: HttpRequest, account_pk: int, team_pk: typing.Optional[int] = None) -> HttpResponse:
        self.perform_account_fetch(request.user, account_pk)
        team = fetch_team_for_account(self.account_fetch_result.account, team_pk)

        form = TeamForm(instance=team)
        if not self.has_permission(AccountPermissionType.ADMINISTER):
            for field in form.fields.values():
                field.disabled = True

        return render(request, "accounts/team_detail.html", self.get_render_context({
            "team": team,
            "account": self.account,
            "form": form
        }))

    def post(self, request: HttpRequest, account_pk: int, team_pk: typing.Optional[int] = None) -> HttpResponse:
        self.perform_account_fetch(request.user, account_pk)

        if not self.has_permission(AccountPermissionType.ADMINISTER):
            raise PermissionDenied

        team = fetch_team_for_account(self.account, team_pk)

        form = TeamForm(request.POST, instance=team)

        if form.is_valid():
            form.save()
            if team_pk is None:
                update_verb = "created"
            else:
                update_verb = "updated"

            messages.success(request, "Team '{}' was {} successfully".format(team.name, update_verb))
            return redirect(reverse('account_team_list', args=(self.account.id,)))

        return render(request, "accounts/team_detail.html", self.get_render_context({
            "team": team,
            "account": self.account,
            "form": form
        }))

class TeamListView(AccountPermissionsMixin, View):
    def get(self, request: HttpRequest, account_pk: int) -> HttpResponse:
        self.perform_account_fetch(request.user, account_pk)

        if not self.account_permissions:
            raise PermissionDenied
        # Assume if they have any Roles for the Account they have access to this page

        return render(request, "accounts/account_teams.html", self.get_render_context({
            "account": self.account,
            "teams": self.account.teams.all
        }))

class TeamMembersView(AccountPermissionsMixin, View):
    def get(self, request: HttpRequest, account_pk: int, team_pk: int) -> HttpResponse:
        self.perform_account_fetch(request.user, account_pk)
        team = fetch_team_for_account(self.account, team_pk)

        return render(request, "accounts/team_members.html", self.get_render_context({
            "account": self.account,
            "team": team
        }))

    def post(self, request: HttpRequest, account_pk: int, team_pk: int) -> HttpResponse:
        self.perform_account_fetch(request.user, account_pk)

        if not self.has_permission(AccountPermissionType.ADMINISTER):
            raise PermissionDenied

        team = fetch_team_for_account(self.account, team_pk)

        action = request.POST.get('action')

        if action in ('add_member', 'remove_member'):
            username = request.POST['name']

            if username:
                user = User.objects.get(username=username)

                if action == 'add_member':
                    if user not in team.members.all():
                        team.members.add(user)
                        messages.success(request, "{} was added to team {}.".format(user.username, team))
                elif action == 'remove_member':
                    if user in team.members.all():
                        team.members.remove(user)
                        messages.success(request, "{} was removed from team {}.".format(user.username, team))

        return redirect(reverse('account_team_members', args=(account_pk, team_pk)))

class TeamProjectsView(AccountPermissionsMixin, View):
    def get(self, request: HttpRequest, account_pk: int, team_pk: int) -> HttpResponse:
        self.perform_account_fetch(request.user, account_pk)
        team = fetch_team_for_account(self.account, team_pk)
        project_roles = ProjectRole.objects.all()
        all_projects = self.account.projects.all()
        existing_project_roles = ProjectAgentRole.objects.filter(project__in=all_projects, agent_id=team.pk,
                                                                 content_type=ContentType.objects.get_for_model(Team))
        assigned_projects = list(map(attrgetter('project'), existing_project_roles))

        unassigned_projects = filter(lambda p: p not in assigned_projects, all_projects)
        # this filtering can't be done by the ORM due to using GenericForeignKey
        # Shouldn't be dealing with millions of projects per account so should be OK

        return render(request, "accounts/team_projects.html", self.get_render_context({
            "account": self.account,
            "team": team,
            "total_project_count": all_projects.count(),
            "existing_project_roles": existing_project_roles,
            "unassigned_projects": list(unassigned_projects),
            "project_roles": project_roles,
            "project_roles_map": json.dumps(
                {existing_project_role.pk: existing_project_role.role_id for existing_project_role in
                 existing_project_roles}),
            "AGENT_ROLE_ID_PREFIX": AGENT_ROLE_ID_PREFIX
        }))

    def post(self, request: HttpRequest, account_pk: int, team_pk: int) -> HttpResponse:
        self.perform_account_fetch(request.user, account_pk)

        if not self.has_permission(AccountPermissionType.ADMINISTER):
            raise PermissionDenied

        team = fetch_team_for_account(self.account, team_pk)

        all_roles = ProjectRole.objects.all()

        role_lookup = {role.pk: role for role in all_roles}

        if request.POST.get('action') == 'add_project':
            role = role_lookup[int(request.POST['role_id'])]
            project = Project.objects.get(pk=request.POST['project_id'])
            ProjectAgentRole.objects.update_or_create({
                'role': role
            }, agent_id=team.pk, content_type=ContentType.objects.get_for_model(Team), project=project)
            messages.success(request, "Project access to {} for {} was saved.".format(project.name, team.name))
        elif request.POST.get('action') == 'remove_project':
            project_agent_role = ProjectAgentRole.objects.get(pk=request.POST['agent_role_id'])
            project = project_agent_role.project
            project_agent_role.delete()
            messages.success(request, "Project access to {} for {} was removed.".format(project.name, team.name))
        elif request.POST.get('action') == 'set_role':
            role = role_lookup[int(request.POST['role_id'])]
            project_agent_role = ProjectAgentRole.objects.get(pk=request.POST['project_agent_role_id'])

            if project_agent_role.agent_type != AgentType.TEAM:
                return JsonResponse({'success': False, 'message': 'The role you trying to update is not for a team.'})

            if project_agent_role.team != team:
                return JsonResponse(
                    {'success': False, 'message': 'The role you trying to update is not for this team.'})

            project_agent_role.role = role
            project_agent_role.save()
            return JsonResponse({'success': True,
                                 'message': "Access to the project '{}' was updated for this team.".format(
                                     project_agent_role.project.get_name())})

        return redirect(reverse("account_team_projects", args=(self.account.pk, team.pk)))
<commit_msg>Catching error and displaying message if username does not exist when adding/removing from team. Ref #157<commit_after>import json
import typing
from operator import attrgetter

from django.contrib import messages
from django.contrib.auth import get_user_model
from django.contrib.contenttypes.models import ContentType
from django.core.exceptions import PermissionDenied
from django.http import HttpRequest, HttpResponse, JsonResponse
from django.shortcuts import render, redirect
from django.urls import reverse
from django.views import View

from accounts.db_facade import fetch_team_for_account
from accounts.forms import TeamForm
from accounts.models import Team, AccountPermissionType
from accounts.views import AccountPermissionsMixin
from projects.permission_models import ProjectRole, ProjectAgentRole, AgentType
from projects.project_models import Project

User = get_user_model()

AGENT_ROLE_ID_PREFIX = 'agent_role_id_'

class TeamDetailView(AccountPermissionsMixin, View):
    def get(self, request: HttpRequest, account_pk: int, team_pk: typing.Optional[int] = None) -> HttpResponse:
        self.perform_account_fetch(request.user, account_pk)
        team = fetch_team_for_account(self.account_fetch_result.account, team_pk)

        form = TeamForm(instance=team)
        if not self.has_permission(AccountPermissionType.ADMINISTER):
            for field in form.fields.values():
                field.disabled = True

        return render(request, "accounts/team_detail.html", self.get_render_context({
            "team": team,
            "account": self.account,
            "form": form
        }))

    def post(self, request: HttpRequest, account_pk: int, team_pk: typing.Optional[int] = None) -> HttpResponse:
        self.perform_account_fetch(request.user, account_pk)

        if not self.has_permission(AccountPermissionType.ADMINISTER):
            raise PermissionDenied

        team = fetch_team_for_account(self.account, team_pk)

        form = TeamForm(request.POST, instance=team)

        if form.is_valid():
            form.save()
            if team_pk is None:
                update_verb = "created"
            else:
                update_verb = "updated"

            messages.success(request, "Team '{}' was {} successfully".format(team.name, update_verb))
            return redirect(reverse('account_team_list', args=(self.account.id,)))

        return render(request, "accounts/team_detail.html", self.get_render_context({
            "team": team,
            "account": self.account,
            "form": form
        }))

class TeamListView(AccountPermissionsMixin, View):
    def get(self, request: HttpRequest, account_pk: int) -> HttpResponse:
        self.perform_account_fetch(request.user, account_pk)

        if not self.account_permissions:
            raise PermissionDenied
        # Assume if they have any Roles for the Account they have access to this page

        return render(request, "accounts/account_teams.html", self.get_render_context({
            "account": self.account,
            "teams": self.account.teams.all
        }))

class TeamMembersView(AccountPermissionsMixin, View):
    def get(self, request: HttpRequest, account_pk: int, team_pk: int) -> HttpResponse:
        self.perform_account_fetch(request.user, account_pk)
        team = fetch_team_for_account(self.account, team_pk)

        return render(request, "accounts/team_members.html", self.get_render_context({
            "account": self.account,
            "team": team
        }))

    def post(self, request: HttpRequest, account_pk: int, team_pk: int) -> HttpResponse:
        self.perform_account_fetch(request.user, account_pk)

        if not self.has_permission(AccountPermissionType.ADMINISTER):
            raise PermissionDenied

        team = fetch_team_for_account(self.account, team_pk)

        action = request.POST.get('action')

        if action in ('add_member', 'remove_member'):
            username = request.POST['name']

            if username:
                try:
                    user = User.objects.get(username=username)
                except User.DoesNotExist:
                    messages.error(request, 'User "{}" does not exist.'.format(username))
                else:
                    if action == 'add_member':
                        if user not in team.members.all():
                            team.members.add(user)
                            messages.success(request, "{} was added to team {}.".format(user.username, team))
                    elif action == 'remove_member':
                        if user in team.members.all():
                            team.members.remove(user)
                            messages.success(request, "{} was removed from team {}.".format(user.username, team))

        return redirect(reverse('account_team_members', args=(account_pk, team_pk)))

class TeamProjectsView(AccountPermissionsMixin, View):
    def get(self, request: HttpRequest, account_pk: int, team_pk: int) -> HttpResponse:
        self.perform_account_fetch(request.user, account_pk)
        team = fetch_team_for_account(self.account, team_pk)
        project_roles = ProjectRole.objects.all()
        all_projects = self.account.projects.all()
        existing_project_roles = ProjectAgentRole.objects.filter(project__in=all_projects, agent_id=team.pk,
                                                                 content_type=ContentType.objects.get_for_model(Team))
        assigned_projects = list(map(attrgetter('project'), existing_project_roles))

        unassigned_projects = filter(lambda p: p not in assigned_projects, all_projects)
        # this filtering can't be done by the ORM due to using GenericForeignKey
        # Shouldn't be dealing with millions of projects per account so should be OK

        return render(request, "accounts/team_projects.html", self.get_render_context({
            "account": self.account,
            "team": team,
            "total_project_count": all_projects.count(),
            "existing_project_roles": existing_project_roles,
            "unassigned_projects": list(unassigned_projects),
            "project_roles": project_roles,
            "project_roles_map": json.dumps(
                {existing_project_role.pk: existing_project_role.role_id for existing_project_role in
                 existing_project_roles}),
            "AGENT_ROLE_ID_PREFIX": AGENT_ROLE_ID_PREFIX
        }))

    def post(self, request: HttpRequest, account_pk: int, team_pk: int) -> HttpResponse:
        self.perform_account_fetch(request.user, account_pk)

        if not self.has_permission(AccountPermissionType.ADMINISTER):
            raise PermissionDenied

        team = fetch_team_for_account(self.account, team_pk)

        all_roles = ProjectRole.objects.all()

        role_lookup = {role.pk: role for role in all_roles}

        if request.POST.get('action') == 'add_project':
            role = role_lookup[int(request.POST['role_id'])]
            project = Project.objects.get(pk=request.POST['project_id'])
            ProjectAgentRole.objects.update_or_create({
                'role': role
            }, agent_id=team.pk, content_type=ContentType.objects.get_for_model(Team), project=project)
            messages.success(request, "Project access to {} for {} was saved.".format(project.name, team.name))
        elif request.POST.get('action') == 'remove_project':
            project_agent_role = ProjectAgentRole.objects.get(pk=request.POST['agent_role_id'])
            project = project_agent_role.project
            project_agent_role.delete()
            messages.success(request, "Project access to {} for {} was removed.".format(project.name, team.name))
        elif request.POST.get('action') == 'set_role':
            role = role_lookup[int(request.POST['role_id'])]
            project_agent_role = ProjectAgentRole.objects.get(pk=request.POST['project_agent_role_id'])

            if project_agent_role.agent_type != AgentType.TEAM:
                return JsonResponse({'success': False, 'message': 'The role you trying to update is not for a team.'})

            if project_agent_role.team != team:
                return JsonResponse(
                    {'success': False, 'message': 'The role you trying to update is not for this team.'})

            project_agent_role.role = role
            project_agent_role.save()
            return JsonResponse({'success': True,
                                 'message': "Access to the project '{}' was updated for this team.".format(
                                     project_agent_role.project.get_name())})

        return redirect(reverse("account_team_projects", args=(self.account.pk, team.pk)))
<|endoftext|>


<commit_before>"""Place of record for the package version"""

__version__ = "2.4.21"
__rpm_version__ = "WILL_BE_SET_BY_RPM_BUILD"
__git_hash__ = "WILL_BE_SET_BY_EGG_BUILD"
<commit_msg>Version bump<commit_after>"""Place of record for the package version"""

__version__ = "2.4.22"
__rpm_version__ = "WILL_BE_SET_BY_RPM_BUILD"
__git_hash__ = "WILL_BE_SET_BY_EGG_BUILD"
<|endoftext|>


<commit_before>"""Place of record for the package version"""

__version__ = "2.0.7"
__rpm_version__ = "WILL_BE_SET_BY_RPM_BUILD"
__git_hash__ = "WILL_BE_SET_BY_EGG_BUILD"
<commit_msg>bump version<commit_after>"""Place of record for the package version"""

__version__ = "2.0.8"
__rpm_version__ = "WILL_BE_SET_BY_RPM_BUILD"
__git_hash__ = "WILL_BE_SET_BY_EGG_BUILD"
<|endoftext|>


<commit_before>#!/usr/bin/env pypy

import sys
import glob

from clickmodels.bootstrap import bootstrap
from clickmodels.inference import *
from clickmodels.input_reader import InputReader

try:
    from config import *
except:
    from clickmodels.config_sample import *

def perpGain(r1, r2):
    """ Perplexity gain of r2 over r1. """
    return (r1 - r2) / (r1 - 1)

def llGain(r1, r2):
    """ Log-likelihood gain of r2 over r1. """
    return (math.exp(r2 - r1) - 1)

def avg(l):
    """ Average of an iterable l. """
    s = 0
    n = 0
    for x in l:
        s += x
        n += 1
    return float(s) / n if n else 0

TESTED_MODEL_PAIRS = ['UBM', 'EB_UBM', 'UBMvsDBN']

if 'RBP' in TESTED_MODEL_PAIRS:
    import scipy
    import scipy.optimize

MODEL_CONSTRUCTORS = {
    'DBN': (lambda config: DbnModel((0.9, 0.9, 0.9, 0.9), config=config),
            lambda config: DbnModel((1.0, 0.9, 1.0, 0.9),
                                    ignoreIntents=False, ignoreLayout=False, config=config)),
    'UBMvsDBN': (lambda config: UbmModel(config=config),
                 lambda config: DbnModel((0.9, 0.9, 0.9, 0.9), config=config)),
    'UBM': (lambda config: UbmModel(config=config),
            lambda config: UbmModel(ignoreIntents=False, ignoreLayout=False, config=config)),
    'EB_UBM': (lambda config: UbmModel(config=config),
               lambda config: EbUbmModel(config=config),
               lambda config: EbUbmModel(ignoreIntents=False, ignoreLayout=False,
                                         config=config)),
    'DCM': (lambda config: DcmModel(config=config),
            lambda config: DcmModel(ignoreIntents=False, ignoreLayout=False, config=config)),
    'RBP': (lambda config: SimplifiedRbpModel(config=config),
            lambda: RbpModel(ignoreIntents=False, ignoreLayout=False, config=config))
}

if __name__ == '__main__':
    if len(sys.argv) != 2:
        print >>sys.stderr, \
            'Usage: {0:s} directory_with_log_files_for_different_days'.format(sys.argv[0])
        sys.exit(1)
    perplexityGains = dict((m, defaultdict(lambda: [])) for m in TESTED_MODEL_PAIRS)
    perplexityGainsPos = [dict((m, defaultdict(lambda: [])) for m in TESTED_MODEL_PAIRS) for pos in xrange(MAX_DOCS_PER_QUERY)]
    llGains = dict((m, defaultdict(lambda: [])) for m in TESTED_MODEL_PAIRS)
    interestingFiles = sorted(glob.glob(sys.argv[1] + '/*'))
    N = len(interestingFiles) // 2
    for fileNumber in xrange(N):
        trainFile = interestingFiles[2 * fileNumber]
        testFile = interestingFiles[2 * fileNumber + 1]
        readInput = InputReader(MIN_DOCS_PER_QUERY, MAX_DOCS_PER_QUERY,
                                EXTENDED_LOG_FORMAT, SERP_SIZE,
                                TRAIN_FOR_METRIC,
                                discard_no_clicks=True)
        trainSessions = readInput(open(trainFile))
        testSessions = readInput(open(testFile))
        config = {
            'MAX_QUERY_ID': readInput.current_query_id + 1,
            'MAX_ITERATIONS': MAX_ITERATIONS,
            'DEBUG': DEBUG,
            'PRETTY_LOG': PRETTY_LOG,
            'MAX_DOCS_PER_QUERY': MAX_DOCS_PER_QUERY,
            'SERP_SIZE': SERP_SIZE,
            'TRANSFORM_LOG': TRANSFORM_LOG,
            'QUERY_INDEPENDENT_PAGER': QUERY_INDEPENDENT_PAGER,
            'DEFAULT_REL': DEFAULT_REL
        }
        for modelName in TESTED_MODEL_PAIRS:
            res = []
            models = MODEL_CONSTRUCTORS[modelName]
            for idx, model in enumerate(models):
                m = model(config)
                m.train(trainSessions)
                currentResult = m.test(testSessions, reportPositionPerplexity=True)
                res.append(currentResult)
                print >>sys.stderr, float(fileNumber) / N, modelName, idx, currentResult
                del m
            for i in xrange(len(models)):
                for j in xrange(i + 1, len(models)):
                    perplexityGains[modelName][(i, j)].append(perpGain(res[i][1], res[j][1]))
                    llGains[modelName][(i, j)].append(llGain(res[i][0], res[j][0]))
                    for pos in xrange(MAX_DOCS_PER_QUERY):
                        perplexityGainsPos[pos][modelName][(i, j)].append(perpGain(res[i][2][pos], res[j][2][pos]))

    for t in ['ll', 'perplexity']:
        print t.upper()
        for m in TESTED_MODEL_PAIRS:
            gainsDict = locals()[t + 'Gains'][m]
            for k, gains in gainsDict.iteritems():
                print m, k, gains, bootstrap(gains)[1]
                if t == 'perplexity':
                    print m, 'POSITION PERPLEXITY GAINS:', k, [[f(perplexityGainsPos[pos][m][k]) for f in [avg, lambda l: bootstrap(l)[1]]] for pos in xrange(MAX_DOCS_PER_QUERY)]

<commit_msg>Remove RBP as it is no longer supported.<commit_after>#!/usr/bin/env pypy

import sys
import glob

from clickmodels.bootstrap import bootstrap
from clickmodels.inference import *
from clickmodels.input_reader import InputReader

try:
    from config import *
except:
    from clickmodels.config_sample import *

def perpGain(r1, r2):
    """ Perplexity gain of r2 over r1. """
    return (r1 - r2) / (r1 - 1)

def llGain(r1, r2):
    """ Log-likelihood gain of r2 over r1. """
    return (math.exp(r2 - r1) - 1)

def avg(l):
    """ Average of an iterable l. """
    s = 0
    n = 0
    for x in l:
        s += x
        n += 1
    return float(s) / n if n else 0

TESTED_MODEL_PAIRS = ['UBM', 'EB_UBM', 'UBMvsDBN']

MODEL_CONSTRUCTORS = {
    'DBN': (lambda config: DbnModel((0.9, 0.9, 0.9, 0.9), config=config),
            lambda config: DbnModel((1.0, 0.9, 1.0, 0.9),
                                    ignoreIntents=False, ignoreLayout=False, config=config)),
    'UBMvsDBN': (lambda config: UbmModel(config=config),
                 lambda config: DbnModel((0.9, 0.9, 0.9, 0.9), config=config)),
    'UBM': (lambda config: UbmModel(config=config),
            lambda config: UbmModel(ignoreIntents=False, ignoreLayout=False, config=config)),
    'EB_UBM': (lambda config: UbmModel(config=config),
               lambda config: EbUbmModel(config=config),
               lambda config: EbUbmModel(ignoreIntents=False, ignoreLayout=False,
                                         config=config)),
    'DCM': (lambda config: DcmModel(config=config),
            lambda config: DcmModel(ignoreIntents=False, ignoreLayout=False, config=config)),
}

if __name__ == '__main__':
    if len(sys.argv) != 2:
        print >>sys.stderr, \
            'Usage: {0:s} directory_with_log_files_for_different_days'.format(sys.argv[0])
        sys.exit(1)
    perplexityGains = dict((m, defaultdict(lambda: [])) for m in TESTED_MODEL_PAIRS)
    perplexityGainsPos = [dict((m, defaultdict(lambda: [])) for m in TESTED_MODEL_PAIRS) for pos in xrange(MAX_DOCS_PER_QUERY)]
    llGains = dict((m, defaultdict(lambda: [])) for m in TESTED_MODEL_PAIRS)
    interestingFiles = sorted(glob.glob(sys.argv[1] + '/*'))
    N = len(interestingFiles) // 2
    for fileNumber in xrange(N):
        trainFile = interestingFiles[2 * fileNumber]
        testFile = interestingFiles[2 * fileNumber + 1]
        readInput = InputReader(MIN_DOCS_PER_QUERY, MAX_DOCS_PER_QUERY,
                                EXTENDED_LOG_FORMAT, SERP_SIZE,
                                TRAIN_FOR_METRIC,
                                discard_no_clicks=True)
        trainSessions = readInput(open(trainFile))
        testSessions = readInput(open(testFile))
        config = {
            'MAX_QUERY_ID': readInput.current_query_id + 1,
            'MAX_ITERATIONS': MAX_ITERATIONS,
            'DEBUG': DEBUG,
            'PRETTY_LOG': PRETTY_LOG,
            'MAX_DOCS_PER_QUERY': MAX_DOCS_PER_QUERY,
            'SERP_SIZE': SERP_SIZE,
            'TRANSFORM_LOG': TRANSFORM_LOG,
            'QUERY_INDEPENDENT_PAGER': QUERY_INDEPENDENT_PAGER,
            'DEFAULT_REL': DEFAULT_REL
        }
        for modelName in TESTED_MODEL_PAIRS:
            res = []
            models = MODEL_CONSTRUCTORS[modelName]
            for idx, model in enumerate(models):
                m = model(config)
                m.train(trainSessions)
                currentResult = m.test(testSessions, reportPositionPerplexity=True)
                res.append(currentResult)
                print >>sys.stderr, float(fileNumber) / N, modelName, idx, currentResult
                del m
            for i in xrange(len(models)):
                for j in xrange(i + 1, len(models)):
                    perplexityGains[modelName][(i, j)].append(perpGain(res[i][1], res[j][1]))
                    llGains[modelName][(i, j)].append(llGain(res[i][0], res[j][0]))
                    for pos in xrange(MAX_DOCS_PER_QUERY):
                        perplexityGainsPos[pos][modelName][(i, j)].append(perpGain(res[i][2][pos], res[j][2][pos]))

    for t in ['ll', 'perplexity']:
        print t.upper()
        for m in TESTED_MODEL_PAIRS:
            gainsDict = locals()[t + 'Gains'][m]
            for k, gains in gainsDict.iteritems():
                print m, k, gains, bootstrap(gains)[1]
                if t == 'perplexity':
                    print m, 'POSITION PERPLEXITY GAINS:', k, [[f(perplexityGainsPos[pos][m][k]) for f in [avg, lambda l: bootstrap(l)[1]]] for pos in xrange(MAX_DOCS_PER_QUERY)]

<|endoftext|>


<commit_before># pylint: disable=invalid-name
# pylint: disable=missing-docstring
# pylint: disable=too-few-public-methods
# pylint: disable=too-many-instance-attributes
# pylint: disable=too-many-locals

import json
import os
import re
import subprocess
import sys
import warnings
import yaml

TESTCASES = list()
INVENTORY = 'test/fixtures/hosts'

HERE = os.path.abspath(os.path.dirname(__file__))
ROLE = re.match(
    r'^.*\/ansible-eos-([^/\s]+)\/test/arista-ansible-role-test$', HERE).group(1)
RUN_CONFIG_BACKUP = '_eos_role_test_{}_running'.format(ROLE)
START_CONFIG_BACKUP = '_eos_role_test_{}_startup'.format(ROLE)

EOS_ROLE_PLAYBOOK = 'test/arista-ansible-role-test/eos_role.yml'
EOS_MODULE_PLAYBOOK = 'test/arista-ansible-role-test/eos_module.yml'

LOG_FILE = '{}/roletest.log'.format(HERE)
try:
    os.remove(LOG_FILE)
except OSError:
    pass
LOG = open(LOG_FILE, 'w')

SEPARATOR = '    ' + '*' * 50

class TestCase(object):
    def __init__(self, **kwargs):
        self.name = kwargs['name']
        self.module = kwargs['module']

        self.host = None

        self.inventory = kwargs.get('inventory')
        self.exitcode = kwargs.get('exitcode', 0)
        self.idempotent = kwargs.get('idempotent', True)
        self.changed = kwargs.get('changed', True)
        self.present = kwargs.get('present')
        self.absent = kwargs.get('absent')

        self.arguments = kwargs.get('arguments', list())
        self.variables = dict()

        # optional properties
        self.setup = kwargs.get('setup', list())
        self.teardown = kwargs.get('teardown', list())

    def __str__(self):
        return self.name

class TestModule(object):
    def __init__(self, testcase):
        self.testcase = testcase
        self.description = 'Test [%s]: %s' % (testcase.module, testcase.name)

    def __call__(self):
        self.output('Run first pass')
        response = self.run_module()
        for device in response:
            hostname = device.keys()[0]
            reported = int(device[hostname]['changed'])
            expected = int(self.testcase.changed)
            msg = ("First pass role execution reported {} task change(s), "
                   "expected {}".format(reported, expected))
            self.output(msg)
            assert reported == expected, msg

        if self.testcase.idempotent:
            self.output('Run second pass')
            response = self.run_module()
            for device in response:
                hostname = device.keys()[0]
                reported = int(device[hostname]['changed'])
                msg = (
                    "Second pass role execution reported {} task change(s), "
                    "expected 0".format(reported))
                self.output(msg)
                assert not reported, msg

        if self.testcase.present:
            desc = 'Validate present configuration'
            self.output(desc)
            response = self.run_validation(self.testcase.present, desc=desc)
            for device in response:
                hostname = device.keys()[0]
                # Result should contain an empty list of updates
                delim = " ---\n"
                updates = device[hostname]['updates']
                msg = ("{} - Expected configuration\n{}{}\n{}not found "
                       "on device '{}'".format(desc, delim,
                                               '\n'.join(updates), delim,
                                               hostname))
                assert device[hostname]['updates'] == [], msg
                # Result should show no changes
                msg = ("{} - Device '{}' reported no updates, but "
                       "returned 'changed'".format(desc, hostname))
                assert device[hostname]['changed'] == False, msg

        if self.testcase.absent:
            desc = 'Validate absent configuration'
            self.output(desc)
            response = self.run_validation(self.testcase.absent, desc=desc)
            for device in response:
                hostname = device.keys()[0]
                # Result should show change has taken place
                msg = (
                    "{} - Entire absent configuration found on device '{}'".
                    format(desc, hostname)
                )
                assert device[hostname]['changed'] == True, msg
                # Compare changes with expected values, sorted at global level
                updates = '\n'.join(device[hostname]['updates'])
                updates = re.split(r'\n(?=\S)', updates)
                updates = '\n'.join(sorted(updates))
                absent = re.split(r'\n(?=\S)', self.testcase.absent.rstrip())
                absent = '\n'.join(sorted(absent))
                msg = ("{} - Some part of absent configuration found "
                       "on device '{}'".format(desc, hostname))
                assert updates == absent, msg

    def setUp(self):
        print("\n{}\n".format(SEPARATOR) +
              "  See run log for complete output:\n  {}".format(LOG_FILE) +
              "\n{}\n".format(SEPARATOR))

        LOG.write("\n\n\n{}\n".format(SEPARATOR) +
                  "  Begin log for {}".format(self.description) +
                  "\n{}\n\n".format(SEPARATOR))

        if self.testcase.setup:
            self.output('Running test case setup commands')
            setup_cmds = self.testcase.setup
            if not isinstance(setup_cmds, list):
                setup_cmds = setup_cmds.splitlines()
            self.output("{}".format(setup_cmds))

            args = {
                'module': 'eos_command',
                'description': 'Run test case setup commands',
                'cmds': ['configure terminal'] + setup_cmds,
            }

            arguments = [json.dumps(args)]

            ret_code, out, err = ansible_playbook(EOS_MODULE_PLAYBOOK,
                                                  arguments=arguments)

            if ret_code != 0:
                LOG.write("Playbook stdout:\n\n{}".format(out))
                LOG.write("Playbook stderr:\n\n{}".format(err))
                raise

    def tearDown(self):
        if self.testcase.teardown:
            self.output('Running test case teardown commands')
            teardown_cmds = self.testcase.teardown
            if not isinstance(teardown_cmds, list):
                teardown_cmds = teardown_cmds.splitlines()
            self.output("{}\n".format(teardown_cmds))

            args = {
                'module': 'eos_command',
                'description': 'Run test case teardown_cmds commands',
                'cmds': ['configure terminal'] + teardown_cmds,
            }

            arguments = [json.dumps(args)]

            ret_code, out, err = ansible_playbook(EOS_MODULE_PLAYBOOK,
                                                  arguments=arguments)

            if ret_code != 0:
                self.output("Playbook stdout:\n\n{}".format(out))
                self.output("Playbook stderr:\n\n{}".format(err))
                warnings.warn("\nError in test case teardown\n\n{}".format(
                    out))

    @classmethod
    def output(cls, text):
        print '>>', str(text)
        LOG.write('++ {}'.format(text) + '\n')

    def run_module(self):
        (retcode, out, _) = self.execute_module()
        out_stripped = re.sub(r'\"config\": \"! Command:.*\\nend\"',
                              '\"config\": \"--- stripped for space ---\"',
                              out)
        LOG.write("PLaybook stdout:\n\n{}".format(out_stripped))
        msg = "Return code: {}, Expected code: {}".format(retcode, self.testcase.exitcode)
        self.output(msg)
        assert retcode == self.testcase.exitcode, msg
        return self.parse_response(out)

    def execute_module(self):
        arguments = [json.dumps(self.testcase.arguments)]
        arguments.append(json.dumps(
            {'rolename': "ansible-eos-{}".format(ROLE)}))
        return ansible_playbook(EOS_ROLE_PLAYBOOK, arguments=arguments)

    def parse_response(self, output, validate=False):
        # Get all the lines after the 'PLAY RECAP ****...' header
        lines = re.sub(r'^.*PLAY RECAP \*+', '', output, 0, re.S).split('\n')
        # Remove any empty lines from the list
        lines = [x for x in lines if x]

        recap = []
        for line in lines:
            match = re.search(r'^(\S+)\s+\:\s+' \
                              r'ok=(\d+)\s+' \
                              r'changed=(\d+)\s+' \
                              r'unreachable=(\d+)\s+' \
                              r'failed=(\d+)', line)
            if not match:
                self.output("Playbook stdout:\n\n{}".format(output))
                raise ValueError("Unable to parse Ansible output for "
                                 "recap information")
            (name, okcount, changed, unreach, failed) = match.groups()
            recap.append({name: {'ok': okcount,
                                 'changed': changed,
                                 'unreachable': unreach,
                                 'failed': failed}})

        if not validate:
            return recap

        updates = []
        for device in recap:
            hostname = device.keys()[0]
            match = re.search(r'(?<!skipping: )\[%s\] => (\{.*\})' % hostname,
                              output, re.M)
            if not match:
                self.output("Playbook stdout:\n\n{}".format(output))
                raise ValueError("Unable to parse Ansible output for "
                                 "result validation")
            result = json.loads(match.group(1))
            updates.append({hostname: result})

        return updates

    def run_validation(self, src, desc='Validate configuration'):
        args = {'module': 'eos_template', 'description': desc, 'src': src, }
        arguments = [json.dumps(args)]
        (ret_code, out, _) = ansible_playbook(EOS_MODULE_PLAYBOOK,
                                              arguments=arguments,
                                              options=['--check'])
        LOG.write(out)
        assert ret_code == 0, "Validation playbook failed execution"
        return self.parse_response(out, validate=True)

def filter_modules(modules, filenames):
    if modules:
        modules = ['{0}.yml'.format(s) for s in modules.split(',')]
        return list(set(modules).intersection(filenames))
    return filenames

def setup():
    print >> sys.stderr, "Test Suite Setup:"

    run_backup = "  Backing up running-config on nodes ..."
    print >> sys.stderr, run_backup
    LOG.write('++ ' + run_backup.strip())
    args = {
        'module': 'eos_command',
        'description': 'Back up running-config on node',
        'cmds': [
            'configure terminal',
            'copy running-config {}'.format(RUN_CONFIG_BACKUP)
        ],
    }
    arguments = [json.dumps(args)]

    ret_code, out, err = ansible_playbook(EOS_MODULE_PLAYBOOK,
                                          arguments=arguments)

    if ret_code != 0:
        LOG.write(str(">> ansible-playbook {} stdout:\n".format(EOS_MODULE_PLAYBOOK), out))
        LOG.write(str(">> ansible-playbook {} stddrr:\n".format(EOS_MODULE_PLAYBOOK), err))
        teardown()
        raise RuntimeError("Error in Test Suite Setup")

    run_backup = "  Backing up startup-config on nodes ..."
    print >> sys.stderr, run_backup
    LOG.write('++ ' + run_backup.strip())
    args = {
        'module': 'eos_command',
        'description': 'Back up startup-config on node',
        'cmds': [
            'configure terminal',
            'copy startup-config {}'.format(START_CONFIG_BACKUP)
        ],
    }
    arguments = [json.dumps(args)]

    ret_code, out, err = ansible_playbook(EOS_MODULE_PLAYBOOK,
                                          arguments=arguments)

    if ret_code != 0:
        LOG.write(str(">> ansible-playbook {} stdout:\n".format(EOS_MODULE_PLAYBOOK), out))
        LOG.write(str(">> ansible-playbook {} stddrr:\n".format(EOS_MODULE_PLAYBOOK), err))
        teardown()
        raise RuntimeError("Error in Test Suite Setup")

    print >> sys.stderr, "  Gathering test cases ..."

    modules = os.environ.get('ANSIBLE_ROLE_TEST_CASES')

    testcases_home = os.path.join(HERE, 'testcases')
    if not os.path.exists(testcases_home):
        print >> sys.stderr, "\n  ***** Testcase directory not found!! *****"
        teardown()
        raise RuntimeError("Testcase path '{}' does not exist".format(testcases_home))

    filenames = os.listdir(testcases_home)

    for module in filter_modules(modules, filenames):
        path = os.path.join(testcases_home, module)
        definition = yaml.load(open(path))

        defaults = definition.get('defaults', {})
        testcases = definition.get('testcases', [])
        if not testcases:
            print >> sys.stderr, ("\n  ***** No testcases defined in "
                                  "module {} *****\n".format(module))
        else:
            for testcase in definition.get('testcases', []):
                kwargs = defaults.copy()
                kwargs.update(testcase)
                TESTCASES.append(TestCase(**kwargs))

    print >> sys.stderr, "  Setup complete\n"

def teardown():
    print >> sys.stderr, "\nTest Suite Teardown:"

    no_teardown = os.environ.get('NO_ANSIBLE_ROLE_TEST_TEARDOWN')

    if no_teardown:
        print >> sys.stderr, ("{}\n"
                              "  Skipping test suite teardown due to "
                              "NO_ANSIBLE_ROLE_TEST_TEARDOWN\n"
                              "  To restore each device to pre-test state "
                              "execute the following commands\n"
                              "  - configure terminal\n"
                              "  - configure replace {}\n"
                              "  - delete {}\n"
                              "  - copy {} startup-config\n"
                              "  - delete {}\n"
                              "{}".format(SEPARATOR, RUN_CONFIG_BACKUP,
                                          RUN_CONFIG_BACKUP,
                                          START_CONFIG_BACKUP,
                                          START_CONFIG_BACKUP, SEPARATOR))
    else:
        # Restore the running-config on the nodes
        # ---------------------------------------
        restore_backup = "  Restoring running-config on nodes ..."
        print >> sys.stderr, restore_backup
        LOG.write('++ ' + restore_backup.strip())
        args = {
            'module': 'eos_command',
            'description': 'Restore running-config from backup',
            'cmds': [
                'configure terminal',
                'configure replace {}'.format(RUN_CONFIG_BACKUP),
                'delete {}'.format(RUN_CONFIG_BACKUP),
            ],
        }
        arguments = [json.dumps(args)]

        # ret_code, out, err = ansible_playbook(CMD_PLAY, arguments=arguments)
        ret_code, out, err = ansible_playbook(EOS_MODULE_PLAYBOOK,
                                              arguments=arguments)

        if ret_code != 0:
            msg = "Error restoring running-config on nodes\n" \
                  "Running ansible-playbook {} -e {}\n" \
                  ">> stdout: {}\n" \
                  ">> stderr: {}\n".format(EOS_MODULE_PLAYBOOK, arguments, out, err)
            warnings.warn(msg)

        # Restore the startup-config on the nodes
        # ---------------------------------------
        restore_backup = "  Restoring startup-config on nodes ..."
        print >> sys.stderr, restore_backup
        LOG.write('++ ' + restore_backup.strip())
        args = {
            'module': 'eos_command',
            'description': 'Restore startup-config from backup',
            'cmds': [
                'configure terminal',
                'copy {} startup-config'.format(START_CONFIG_BACKUP),
                'delete {}'.format(START_CONFIG_BACKUP),
            ],
        }
        arguments = [json.dumps(args)]

        # ret_code, out, err = ansible_playbook(CMD_PLAY, arguments=arguments)
        ret_code, out, err = ansible_playbook(EOS_MODULE_PLAYBOOK,
                                              arguments=arguments)

        if ret_code != 0:
            msg = "Error restoring startup-config on nodes\n" \
                  "Running ansible-playbook {} -e {}\n" \
                  ">> stdout: {}\n" \
                  ">> stderr: {}\n".format(EOS_MODULE_PLAYBOOK, arguments, out, err)
            warnings.warn(msg)

    print >> sys.stderr, "  Teardown complete"

def test_module():
    for testcase in TESTCASES:
        yield TestModule(testcase)

def ansible_playbook(playbook, arguments=None, options=None):
    if arguments is None:
        arguments = []
    if options is None:
        options = []

    command = ['ansible-playbook']

    command.append(playbook)
    command.extend(['-i', INVENTORY])
    for arg in arguments:
        command.extend(['-e', arg])
    for opt in options:
        command.append(opt)
    command.append('-vvv')

    # Format the command string for output on error - for easier
    # copy/paste for manual run
    cmdstr = ''
    for segment in command:
        if segment[0] == '{':
            cmdstr = cmdstr + "\'{}\' ".format(segment)
        else:
            cmdstr = cmdstr + "{} ".format(segment)
    LOG.write("-- Ansible playbook command:\n-- {}\n".format(cmdstr))

    stdout = subprocess.PIPE
    stderr = subprocess.PIPE
    proc = subprocess.Popen(command, stdout=stdout, stderr=stderr)
    out, err = proc.communicate()

    return (proc.returncode, out, err)
<commit_msg>Pull in arista-ansible-role-test updates<commit_after># pylint: disable=invalid-name
# pylint: disable=missing-docstring
# pylint: disable=too-few-public-methods
# pylint: disable=too-many-instance-attributes
# pylint: disable=too-many-locals

import json
import os
import re
import subprocess
import sys
import warnings
import yaml

TESTCASES = list()
INVENTORY = 'test/fixtures/hosts'

HERE = os.path.abspath(os.path.dirname(__file__))
ROLE = re.match(
    r'^.*\/ansible-eos-([^/\s]+)\/test/arista-ansible-role-test$', HERE).group(1)
RUN_CONFIG_BACKUP = '_eos_role_test_{}_running'.format(ROLE)
START_CONFIG_BACKUP = '_eos_role_test_{}_startup'.format(ROLE)

EOS_ROLE_PLAYBOOK = 'test/arista-ansible-role-test/eos_role.yml'
EOS_MODULE_PLAYBOOK = 'test/arista-ansible-role-test/eos_module.yml'

LOG_FILE = '{}/roletest.log'.format(HERE)
try:
    os.remove(LOG_FILE)
except OSError:
    pass
LOG = open(LOG_FILE, 'w')

SEPARATOR = '    ' + '*' * 50

class TestCase(object):
    def __init__(self, **kwargs):
        self.name = kwargs['name']
        self.module = kwargs['module']

        self.host = None

        self.inventory = kwargs.get('inventory')
        self.exitcode = kwargs.get('exitcode', 0)
        self.idempotent = kwargs.get('idempotent', True)
        self.changed = kwargs.get('changed', True)
        self.present = kwargs.get('present')
        self.absent = kwargs.get('absent')

        self.arguments = kwargs.get('arguments', list())
        self.variables = dict()

        # optional properties
        self.setup = kwargs.get('setup', list())
        self.teardown = kwargs.get('teardown', list())

    def __str__(self):
        return self.name

class TestModule(object):
    def __init__(self, testcase):
        self.testcase = testcase
        self.description = 'Test [%s]: %s' % (testcase.module, testcase.name)

    def __call__(self):
        self.output('Run first pass')
        response = self.run_module()
        for device in response:
            hostname = device.keys()[0]
            reported = int(device[hostname]['changed'])
            expected = int(self.testcase.changed)
            msg = ("First pass role execution reported {} task change(s), "
                   "expected {}".format(reported, expected))
            self.output(msg)
            assert reported == expected, msg

        if self.testcase.idempotent:
            self.output('Run second pass')
            response = self.run_module()
            for device in response:
                hostname = device.keys()[0]
                reported = int(device[hostname]['changed'])
                msg = (
                    "Second pass role execution reported {} task change(s), "
                    "expected 0".format(reported))
                self.output(msg)
                assert not reported, msg

        if self.testcase.present:
            desc = 'Validate present configuration'
            self.output(desc)
            response = self.run_validation(self.testcase.present, desc=desc)
            for device in response:
                hostname = device.keys()[0]
                # Result should contain an empty list of updates
                delim = " ---\n"
                updates = device[hostname]['updates']
                msg = ("{} - Expected configuration\n{}{}\n{}not found "
                       "on device '{}'".format(desc, delim,
                                               '\n'.join(updates), delim,
                                               hostname))
                assert device[hostname]['updates'] == [], msg
                # Result should show no changes
                msg = ("{} - Device '{}' reported no updates, but "
                       "returned 'changed'".format(desc, hostname))
                assert device[hostname]['changed'] == False, msg

        if self.testcase.absent:
            desc = 'Validate absent configuration'
            self.output(desc)
            response = self.run_validation(self.testcase.absent, desc=desc)
            for device in response:
                hostname = device.keys()[0]
                # Result should show change has taken place
                msg = (
                    "{} - Entire absent configuration found on device '{}'".
                    format(desc, hostname)
                )
                assert device[hostname]['changed'] == True, msg
                # Compare changes with expected values, sorted at global level
                updates = '\n'.join(device[hostname]['updates'])
                updates = re.split(r'\n(?=\S)', updates)
                updates = '\n'.join(sorted(updates))
                absent = re.split(r'\n(?=\S)', self.testcase.absent.rstrip())
                absent = '\n'.join(sorted(absent))
                msg = ("{} - Some part of absent configuration found "
                       "on device '{}'".format(desc, hostname))
                assert updates == absent, msg

    def setUp(self):
        print("\n{}\n".format(SEPARATOR) +
              "  See run log for complete output:\n  {}".format(LOG_FILE) +
              "\n{}\n".format(SEPARATOR))

        LOG.write("\n\n\n{}\n".format(SEPARATOR) +
                  "  Begin log for {}".format(self.description) +
                  "\n{}\n\n".format(SEPARATOR))

        if self.testcase.setup:
            self.output('Running test case setup commands')
            setup_cmds = self.testcase.setup
            if not isinstance(setup_cmds, list):
                setup_cmds = setup_cmds.splitlines()
            self.output("{}".format(setup_cmds))

            args = {
                'module': 'eos_command',
                'description': 'Run test case setup commands',
                'cmds': ['configure terminal'] + setup_cmds,
            }

            arguments = [json.dumps(args)]

            ret_code, out, err = ansible_playbook(EOS_MODULE_PLAYBOOK,
                                                  arguments=arguments)

            if ret_code != 0:
                LOG.write("Playbook stdout:\n\n{}".format(out))
                LOG.write("Playbook stderr:\n\n{}".format(err))
                raise

    def tearDown(self):
        if self.testcase.teardown:
            self.output('Running test case teardown commands')
            teardown_cmds = self.testcase.teardown
            if not isinstance(teardown_cmds, list):
                teardown_cmds = teardown_cmds.splitlines()
            self.output("{}\n".format(teardown_cmds))

            args = {
                'module': 'eos_command',
                'description': 'Run test case teardown_cmds commands',
                'cmds': ['configure terminal'] + teardown_cmds,
            }

            arguments = [json.dumps(args)]

            ret_code, out, err = ansible_playbook(EOS_MODULE_PLAYBOOK,
                                                  arguments=arguments)

            if ret_code != 0:
                self.output("Playbook stdout:\n\n{}".format(out))
                self.output("Playbook stderr:\n\n{}".format(err))
                warnings.warn("\nError in test case teardown\n\n{}".format(
                    out))

    @classmethod
    def output(cls, text):
        print '>>', str(text)
        LOG.write('++ {}'.format(text) + '\n')

    def run_module(self):
        (retcode, out, _) = self.execute_module()
        out_stripped = re.sub(r'\"config\": \"! Command:.*\\nend\"',
                              '\"config\": \"--- stripped for space ---\"',
                              out)
        LOG.write("PLaybook stdout:\n\n{}".format(out_stripped))
        msg = "Return code: {}, Expected code: {}".format(retcode, self.testcase.exitcode)
        self.output(msg)
        assert retcode == self.testcase.exitcode, msg
        return self.parse_response(out)

    def execute_module(self):
        arguments = [json.dumps(self.testcase.arguments)]
        arguments.append(json.dumps(
            {'rolename': "ansible-eos-{}".format(ROLE)}))
        return ansible_playbook(EOS_ROLE_PLAYBOOK, arguments=arguments)

    def parse_response(self, output, validate=False):
        # Get all the lines after the 'PLAY RECAP ****...' header
        lines = re.sub(r'^.*PLAY RECAP \*+', '', output, 0, re.S).split('\n')
        # Remove any empty lines from the list
        lines = [x for x in lines if x]

        recap = []
        for line in lines:
            match = re.search(r'^(\S+)\s+\:\s+' \
                              r'ok=(\d+)\s+' \
                              r'changed=(\d+)\s+' \
                              r'unreachable=(\d+)\s+' \
                              r'failed=(\d+)', line)
            if not match:
                self.output("Playbook stdout:\n\n{}".format(output))
                raise ValueError("Unable to parse Ansible output for "
                                 "recap information")
            (name, okcount, changed, unreach, failed) = match.groups()
            recap.append({name: {'ok': okcount,
                                 'changed': changed,
                                 'unreachable': unreach,
                                 'failed': failed}})

        if not validate:
            return recap

        updates = []
        for device in recap:
            hostname = device.keys()[0]
            match = re.search(r'(?<!skipping: )\[%s\] => (\{.*\})' % hostname,
                              output, re.M)
            if not match:
                self.output("Playbook stdout:\n\n{}".format(output))
                raise ValueError("Unable to parse Ansible output for "
                                 "result validation")
            result = json.loads(match.group(1))
            updates.append({hostname: result})

        return updates

    def run_validation(self, src, desc='Validate configuration'):
        args = {'module': 'eos_template', 'description': desc, 'src': src, }
        arguments = [json.dumps(args)]
        (ret_code, out, _) = ansible_playbook(EOS_MODULE_PLAYBOOK,
                                              arguments=arguments,
                                              options=['--check'])
        LOG.write(out)
        assert ret_code == 0, "Validation playbook failed execution"
        return self.parse_response(out, validate=True)

def filter_modules(modules, filenames):
    if modules:
        modules = ['{0}.yml'.format(s) for s in modules.split(',')]
        return list(set(modules).intersection(filenames))
    return filenames

def setup():
    print >> sys.stderr, "Test Suite Setup:"

    run_backup = "  Backing up running-config on nodes ..."
    print >> sys.stderr, run_backup
    LOG.write('++ ' + run_backup.strip())
    args = {
        'module': 'eos_command',
        'description': 'Back up running-config on node',
        'cmds': [
            'configure terminal',
            'copy running-config {}'.format(RUN_CONFIG_BACKUP)
        ],
    }
    arguments = [json.dumps(args)]

    ret_code, out, err = ansible_playbook(EOS_MODULE_PLAYBOOK,
                                          arguments=arguments)

    if ret_code != 0:
        LOG.write(">> ansible-playbook {} stdout:\n{}".format(EOS_MODULE_PLAYBOOK, out))
        LOG.write(">> ansible-playbook {} stddrr:\n{}".format(EOS_MODULE_PLAYBOOK, err))
        teardown()
        raise RuntimeError("Error in Test Suite Setup")

    run_backup = "  Backing up startup-config on nodes ..."
    print >> sys.stderr, run_backup
    LOG.write('++ ' + run_backup.strip())
    args = {
        'module': 'eos_command',
        'description': 'Back up startup-config on node',
        'cmds': [
            'configure terminal',
            'copy startup-config {}'.format(START_CONFIG_BACKUP)
        ],
    }
    arguments = [json.dumps(args)]

    ret_code, out, err = ansible_playbook(EOS_MODULE_PLAYBOOK,
                                          arguments=arguments)

    if ret_code != 0:
        LOG.write(">> ansible-playbook {} stdout:\n{}".format(EOS_MODULE_PLAYBOOK, out))
        LOG.write(">> ansible-playbook {} stddrr:\n{}".format(EOS_MODULE_PLAYBOOK, err))
        teardown()
        raise RuntimeError("Error in Test Suite Setup")

    print >> sys.stderr, "  Gathering test cases ..."

    modules = os.environ.get('ANSIBLE_ROLE_TEST_CASES')

    testcases_home = os.path.join(HERE, 'testcases')
    if not os.path.exists(testcases_home):
        print >> sys.stderr, "\n  ***** Testcase directory not found!! *****"
        teardown()
        raise RuntimeError("Testcase path '{}' does not exist".format(testcases_home))

    filenames = os.listdir(testcases_home)

    for module in filter_modules(modules, filenames):
        path = os.path.join(testcases_home, module)
        definition = yaml.load(open(path))

        defaults = definition.get('defaults', {})
        testcases = definition.get('testcases', [])
        if not testcases:
            print >> sys.stderr, ("\n  ***** No testcases defined in "
                                  "module {} *****\n".format(module))
        else:
            for testcase in definition.get('testcases', []):
                kwargs = defaults.copy()
                kwargs.update(testcase)
                TESTCASES.append(TestCase(**kwargs))

    print >> sys.stderr, "  Setup complete\n"

def teardown():
    print >> sys.stderr, "\nTest Suite Teardown:"

    no_teardown = os.environ.get('NO_ANSIBLE_ROLE_TEST_TEARDOWN')

    if no_teardown:
        print >> sys.stderr, ("{}\n"
                              "  Skipping test suite teardown due to "
                              "NO_ANSIBLE_ROLE_TEST_TEARDOWN\n"
                              "  To restore each device to pre-test state "
                              "execute the following commands\n"
                              "  - configure terminal\n"
                              "  - configure replace {}\n"
                              "  - delete {}\n"
                              "  - copy {} startup-config\n"
                              "  - delete {}\n"
                              "{}".format(SEPARATOR, RUN_CONFIG_BACKUP,
                                          RUN_CONFIG_BACKUP,
                                          START_CONFIG_BACKUP,
                                          START_CONFIG_BACKUP, SEPARATOR))
    else:
        # Restore the running-config on the nodes
        # ---------------------------------------
        restore_backup = "  Restoring running-config on nodes ..."
        print >> sys.stderr, restore_backup
        LOG.write('++ ' + restore_backup.strip())
        args = {
            'module': 'eos_command',
            'description': 'Restore running-config from backup',
            'cmds': [
                'configure terminal',
                'configure replace {}'.format(RUN_CONFIG_BACKUP),
                'delete {}'.format(RUN_CONFIG_BACKUP),
            ],
        }
        arguments = [json.dumps(args)]

        # ret_code, out, err = ansible_playbook(CMD_PLAY, arguments=arguments)
        ret_code, out, err = ansible_playbook(EOS_MODULE_PLAYBOOK,
                                              arguments=arguments)

        if ret_code != 0:
            msg = "Error restoring running-config on nodes\n" \
                  "Running ansible-playbook {} -e {}\n" \
                  ">> stdout: {}\n" \
                  ">> stderr: {}\n".format(EOS_MODULE_PLAYBOOK, arguments, out, err)
            warnings.warn(msg)

        # Restore the startup-config on the nodes
        # ---------------------------------------
        restore_backup = "  Restoring startup-config on nodes ..."
        print >> sys.stderr, restore_backup
        LOG.write('++ ' + restore_backup.strip())
        args = {
            'module': 'eos_command',
            'description': 'Restore startup-config from backup',
            'cmds': [
                'configure terminal',
                'copy {} startup-config'.format(START_CONFIG_BACKUP),
                'delete {}'.format(START_CONFIG_BACKUP),
            ],
        }
        arguments = [json.dumps(args)]

        # ret_code, out, err = ansible_playbook(CMD_PLAY, arguments=arguments)
        ret_code, out, err = ansible_playbook(EOS_MODULE_PLAYBOOK,
                                              arguments=arguments)

        if ret_code != 0:
            msg = "Error restoring startup-config on nodes\n" \
                  "Running ansible-playbook {} -e {}\n" \
                  ">> stdout: {}\n" \
                  ">> stderr: {}\n".format(EOS_MODULE_PLAYBOOK, arguments, out, err)
            warnings.warn(msg)

    print >> sys.stderr, "  Teardown complete"

def test_module():
    for testcase in TESTCASES:
        yield TestModule(testcase)

def ansible_playbook(playbook, arguments=None, options=None):
    if arguments is None:
        arguments = []
    if options is None:
        options = []

    command = ['ansible-playbook']

    command.append(playbook)
    command.extend(['-i', INVENTORY])
    for arg in arguments:
        command.extend(['-e', arg])
    for opt in options:
        command.append(opt)
    command.append('-vvv')

    # Format the command string for output on error - for easier
    # copy/paste for manual run
    cmdstr = ''
    for segment in command:
        if segment[0] == '{':
            cmdstr = cmdstr + "\'{}\' ".format(segment)
        else:
            cmdstr = cmdstr + "{} ".format(segment)
    LOG.write("-- Ansible playbook command:\n-- {}\n".format(cmdstr))

    stdout = subprocess.PIPE
    stderr = subprocess.PIPE
    proc = subprocess.Popen(command, stdout=stdout, stderr=stderr)
    out, err = proc.communicate()

    return (proc.returncode, out, err)
<|endoftext|>


<commit_before>import logging
from logging import Handler

try:
    import sqlalchemy as sa
    from sqlalchemy import Table, Column, Text, Integer, Float, Boolean
except ImportError:
    pass

# TODO: expand to full set of info
def create_workflows_table(meta):
    return Table(
            'workflows', meta,
            Column('task_run_id', Text, nullable=False, primary_key=True),
            Column('time_began', Text, nullable=False),
            Column('time_completed', Text),
            # Column('host', Text, nullable=False),
            # Column('user', Text, nullable=False),
            Column('rundir', Text, nullable=False)
    )

# TODO: expand to full set of info
def create_task_status_table(task_id, run_id, meta):
    table_name = run_id + str(task_id)
    return Table(
          table_name, meta,
          Column('task_id', Integer, sa.ForeignKey(run_id + '.task_id'), nullable=False),
          Column('task_status', Integer, nullable=False),
          Column('task_status_name', Integer, nullable=False),
          Column('timestamp', Text, nullable=False, primary_key=True),
          Column('task_run_id', Text, sa.ForeignKey('workflows.task_run_id'), nullable=False),
          Column('task_fail_count', Integer, nullable=False),
          Column('task_fail_history', Text, nullable=True),
    )

# TODO: expand to full set of info
def create_workflow_table(run_id, meta):
    table_name = run_id
    return Table(
          table_name, meta,
          Column('task_id', Integer, primary_key=True, nullable=False),
          Column('task_run_id', Text, sa.ForeignKey('workflows.task_run_id'), nullable=False),
          Column('task_executor', Text, nullable=False),
          Column('task_fn_hash', Text, nullable=False),
          Column('task_time_started', Text, nullable=False),
          Column('task_time_completed', Text, nullable=True),
          Column('task_memoize', Boolean, nullable=False),
          # Column('task_kwargs.input', Text, nullable=True),
          # Column('task_kwargs.output', Text, nullable=True),
          # Column('task_kwargs.stdin', Text, nullable=True),
          # Column('task_kwargs.stdout', Text, nullable=True),
    )

# TODO: expand to full set of info
def create_task_resource_table(task_id, run_id, meta):
    table_name = run_id + str(task_id)
    return Table(
          table_name + '_resources', meta,
          Column('task_id', Integer, sa.ForeignKey(run_id + '.task_id'), nullable=False),
          Column('timestamp', Text, nullable=False, primary_key=True),
          Column('task_run_id', Text, sa.ForeignKey('workflows.task_run_id'), nullable=False),
          Column('psutil_process_pid', Integer, nullable=True),
          Column('psutil_process_cpu_percent', Float, nullable=True),
          Column('psutil_process_memory_percent', Float, nullable=True),
          Column('psutil_process_children_count', Integer, nullable=True),
          Column('psutil_process_time_user', Float, nullable=True),
          Column('psutil_process_time_system', Float, nullable=True),
          Column('psutil_process_memory_virtual', Float, nullable=True),
          Column('psutil_process_memory_resident', Float, nullable=True),
          Column('psutil_process_disk_read', Float, nullable=True),
          Column('psutil_process_disk_write', Float, nullable=True),
          Column('psutil_process_status', Text, nullable=True),
    )

class DatabaseHandler(Handler):
    def __init__(self, elink):
        logging.Handler.__init__(self)
        self.eng = sa.create_engine(elink)

    def emit(self, record):
        with self.eng.connect() as con:
            meta = sa.MetaData()
            # Having what i think is race conditions so try a couple times and don't complain if breaks
            trys = 3
            for t in range(trys):
                failed = False
                try:
                    meta.reflect(bind=con)
                except Exception:
                    print('reflect failed on attempt', t)
                    failed = True
                if failed:
                    pass
                else:
                    print('successful reflect on attempt', t)
                    break
            else:
                print('failed reflect and not logging after', trys, 'attempts')
                return

            info = {key: value for key, value in record.__dict__.items() if not key.startswith("__")}
            # formating values to convert from python or parsl to db standards
            info['task_fail_history'] = ', '.join([str(h) for h in info['task_fail_history']]) if info.get('task_fail_history', None) is not None else None
            info['timestamp'] = record.created
            run_id = info['task_run_id']
            # create workflows table if this is a new database without one
            if 'workflows' not in meta.tables.keys():
                workflows = create_workflows_table(meta)
                workflows.create(con, checkfirst=True)
            # if this is the first sight of the workflow, add it to the workflows table
            if len(con.execute(meta.tables['workflows'].select(meta.tables['workflows'].c.task_run_id == run_id)).fetchall()) == 0:
                try:
                    workflows = meta.tables['workflows']
                    ins = workflows.insert().values(**{k: v for k, v in info.items() if k in workflows.c})
                    con.execute(ins)
                    print(run_id + " was added to the workflows table")
                except sa.exc.IntegrityError as e:
                    print(e)
                    print(dir(e))

            # create workflow table if this is a new run without one
            if run_id not in meta.tables.keys():
                workflow = create_workflow_table(run_id, meta)
                workflow.create(con, checkfirst=True)

            # check to make sure it is a task log and not just a workflow overview log
            if info.get('task_id', None) is not None:
                if 'psutil_process_cpu_percent' in info.keys():
                    # TODO: only use this if it is a task resource update and not a task status update
                    # if this is a task resource update then handle that, if the resource table DNE then create it
                    if (run_id + str(info['task_id']) + "_resources") not in meta.tables.keys():
                        task_resource_table = create_task_resource_table(info['task_id'], run_id, meta)
                        task_resource_table.create(con, checkfirst=True)
                        con.execute(task_resource_table.insert().values(**{k: v for k, v in info.items() if k in task_resource_table.c}))
                        print(task_resource_table, 'table was created and had a task resource update added')
                    # if this resource table already exists, just insert the update
                    else:
                        task_resource_table = meta.tables[run_id + str(info['task_id']) + '_resources']
                        con.execute(task_resource_table.insert().values(**{k: v for k, v in info.items() if k in task_resource_table.c}))
                        print(task_resource_table, 'had a task resource update added')
                    return

                # if this is the first sight of the task in the workflow, add it to the workflow table
                if len(con.execute(meta.tables[run_id].select(meta.tables[run_id].c.task_id == info['task_id'])).fetchall()) == 0:
                    workflow = meta.tables[run_id]
                    ins = workflow.insert().values(**{k: v for k, v in info.items() if k in workflow.c})
                    con.execute(ins)
                    print('Task ' + str(info['task_id']) + " was added to the workflow table")

                if 'task_status' in info.keys():
                    # TODO: only fire this if it is a task status update and not a task resource update
                    # if this is the first sight of a task, create a task_status_table to hold this task's updates
                    if (run_id + str(info['task_id'])) not in meta.tables.keys():
                        task_status_table = create_task_status_table(info['task_id'], run_id, meta)
                        task_status_table.create(con, checkfirst=True)
                        con.execute(task_status_table.insert().values(**{k: v for k, v in info.items() if k in task_status_table.c}))
                        print(task_status_table, 'table was created and had a task status update added')
                    # if this status table already exists, just insert the update
                    else:
                        task_status_table = meta.tables[run_id + str(info['task_id'])]
                        con.execute(task_status_table.insert().values(**{k: v for k, v in info.items() if k in task_status_table.c}))
                        print(task_status_table, 'had a task status update added')
                    return
<commit_msg>Now updates recorde with time_completed.<commit_after>import logging
from logging import Handler

try:
    import sqlalchemy as sa
    from sqlalchemy import Table, Column, Text, Integer, Float, Boolean
except ImportError:
    pass

# TODO: expand to full set of info
def create_workflows_table(meta):
    return Table(
            'workflows', meta,
            Column('task_run_id', Text, nullable=False, primary_key=True),
            Column('time_began', Text, nullable=False),
            Column('time_completed', Text),
            # Column('host', Text, nullable=False),
            # Column('user', Text, nullable=False),
            Column('rundir', Text, nullable=False)
    )

# TODO: expand to full set of info
def create_task_status_table(task_id, run_id, meta):
    table_name = run_id + str(task_id)
    return Table(
          table_name, meta,
          Column('task_id', Integer, sa.ForeignKey(run_id + '.task_id'), nullable=False),
          Column('task_status', Integer, nullable=False),
          Column('task_status_name', Integer, nullable=False),
          Column('timestamp', Text, nullable=False, primary_key=True),
          Column('task_run_id', Text, sa.ForeignKey('workflows.task_run_id'), nullable=False),
          Column('task_fail_count', Integer, nullable=False),
          Column('task_fail_history', Text, nullable=True),
    )

# TODO: expand to full set of info
def create_workflow_table(run_id, meta):
    table_name = run_id
    return Table(
          table_name, meta,
          Column('task_id', Integer, primary_key=True, nullable=False),
          Column('task_run_id', Text, sa.ForeignKey('workflows.task_run_id'), nullable=False),
          Column('task_executor', Text, nullable=False),
          Column('task_fn_hash', Text, nullable=False),
          Column('task_time_started', Text, nullable=False),
          Column('task_time_completed', Text, nullable=True),
          Column('task_memoize', Boolean, nullable=False),
          # Column('task_kwargs.input', Text, nullable=True),
          # Column('task_kwargs.output', Text, nullable=True),
          # Column('task_kwargs.stdin', Text, nullable=True),
          # Column('task_kwargs.stdout', Text, nullable=True),
    )

# TODO: expand to full set of info
def create_task_resource_table(task_id, run_id, meta):
    table_name = run_id + str(task_id)
    return Table(
          table_name + '_resources', meta,
          Column('task_id', Integer, sa.ForeignKey(run_id + '.task_id'), nullable=False),
          Column('timestamp', Text, nullable=False, primary_key=True),
          Column('task_run_id', Text, sa.ForeignKey('workflows.task_run_id'), nullable=False),
          Column('psutil_process_pid', Integer, nullable=True),
          Column('psutil_process_cpu_percent', Float, nullable=True),
          Column('psutil_process_memory_percent', Float, nullable=True),
          Column('psutil_process_children_count', Integer, nullable=True),
          Column('psutil_process_time_user', Float, nullable=True),
          Column('psutil_process_time_system', Float, nullable=True),
          Column('psutil_process_memory_virtual', Float, nullable=True),
          Column('psutil_process_memory_resident', Float, nullable=True),
          Column('psutil_process_disk_read', Float, nullable=True),
          Column('psutil_process_disk_write', Float, nullable=True),
          Column('psutil_process_status', Text, nullable=True),
    )

class DatabaseHandler(Handler):
    def __init__(self, elink):
        logging.Handler.__init__(self)
        self.eng = sa.create_engine(elink)

    def emit(self, record):
        with self.eng.connect() as con:
            meta = sa.MetaData()
            # Having what i think is an issue to reflect so try a couple times and don't complain if breaks
            trys = 3
            for t in range(trys):
                failed = False
                try:
                    meta.reflect(bind=con)
                except Exception:
                    print('reflect failed on attempt', t)
                    failed = True
                if failed:
                    pass
                else:
                    print('successful reflect on attempt', t)
                    break
            else:
                print('failed reflect and not logging after', trys, 'attempts')
                return

            info = {key: value for key, value in record.__dict__.items() if not key.startswith("__")}
            # formating values to convert from python or parsl to db standards
            info['task_fail_history'] = ', '.join([str(h) for h in info['task_fail_history']]) if info.get('task_fail_history', None) is not None else None
            info['timestamp'] = record.created
            run_id = info['task_run_id']

            if 'time_completed' in info.keys() and info['time_completed'] != 'None':
                workflows = meta.tables['workflows']
                up = workflows.update().values(time_completed=info['time_completed']).where(workflows.c.task_run_id == run_id)
                con.execute(up)
                return
            if 'task_time_completed' in info.keys() and info['task_time_completed'] is not None:
                workflow = meta.tables[run_id]
                up = workflow.update().values(task_time_completed=info['task_time_completed']).where(workflow.c.task_id == info['task_id'])
                con.execute(up)
                return

            # create workflows table if this is a new database without one
            if 'workflows' not in meta.tables.keys():
                workflows = create_workflows_table(meta)
                workflows.create(con, checkfirst=True)
            # if this is the first sight of the workflow, add it to the workflows table
            if len(con.execute(meta.tables['workflows'].select(meta.tables['workflows'].c.task_run_id == run_id)).fetchall()) == 0:
                try:
                    workflows = meta.tables['workflows']
                    ins = workflows.insert().values(**{k: v for k, v in info.items() if k in workflows.c})
                    con.execute(ins)
                    print(run_id + " was added to the workflows table")
                except sa.exc.IntegrityError as e:
                    print(e)
                    print(dir(e))

            # create workflow table if this is a new run without one
            if run_id not in meta.tables.keys():
                workflow = create_workflow_table(run_id, meta)
                workflow.create(con, checkfirst=True)

            # check to make sure it is a task log and not just a workflow overview log
            if info.get('task_id', None) is not None:
                if 'psutil_process_cpu_percent' in info.keys():
                    # TODO: only use this if it is a task resource update and not a task status update
                    # if this is a task resource update then handle that, if the resource table DNE then create it
                    if (run_id + str(info['task_id']) + "_resources") not in meta.tables.keys():
                        task_resource_table = create_task_resource_table(info['task_id'], run_id, meta)
                        task_resource_table.create(con, checkfirst=True)
                        con.execute(task_resource_table.insert().values(**{k: v for k, v in info.items() if k in task_resource_table.c}))
                        print(task_resource_table, 'table was created and had a task resource update added')
                    # if this resource table already exists, just insert the update
                    else:
                        task_resource_table = meta.tables[run_id + str(info['task_id']) + '_resources']
                        con.execute(task_resource_table.insert().values(**{k: v for k, v in info.items() if k in task_resource_table.c}))
                        print(task_resource_table, 'had a task resource update added')
                    return

                # if this is the first sight of the task in the workflow, add it to the workflow table
                if len(con.execute(meta.tables[run_id].select(meta.tables[run_id].c.task_id == info['task_id'])).fetchall()) == 0:
                    workflow = meta.tables[run_id]
                    ins = workflow.insert().values(**{k: v for k, v in info.items() if k in workflow.c})
                    con.execute(ins)
                    print('Task ' + str(info['task_id']) + " was added to the workflow table")

                if 'task_status' in info.keys():
                    # TODO: only fire this if it is a task status update and not a task resource update
                    # if this is the first sight of a task, create a task_status_table to hold this task's updates
                    if (run_id + str(info['task_id'])) not in meta.tables.keys():
                        task_status_table = create_task_status_table(info['task_id'], run_id, meta)
                        task_status_table.create(con, checkfirst=True)
                        con.execute(task_status_table.insert().values(**{k: v for k, v in info.items() if k in task_status_table.c}))
                        print(task_status_table, 'table was created and had a task status update added')
                    # if this status table already exists, just insert the update
                    else:
                        task_status_table = meta.tables[run_id + str(info['task_id'])]
                        con.execute(task_status_table.insert().values(**{k: v for k, v in info.items() if k in task_status_table.c}))
                        print(task_status_table, 'had a task status update added')
                    return
<|endoftext|>


<commit_before>__author__="UShareSoft"
__date__ ="$Jan 21, 2014 5:48:54 PM$"

import argparse

class ArgumentParserError(Exception): pass

class ArgumentParser(argparse.ArgumentParser):

	def error(self, message):
        	raise ArgumentParserError(message)

        def format_help(self):
                formatter = self._get_formatter()

                # description
                formatter.add_text(self.description)

                # usage
                formatter.add_usage(self.usage, self._actions,
                                    self._mutually_exclusive_groups)

                # positionals, optionals and user-defined groups
                for action_group in self._action_groups:
                    formatter.start_section(action_group.title)
                    formatter.add_text(action_group.description)
                    formatter.add_arguments(action_group._group_actions)
                    formatter.end_section()

                # epilog
                formatter.add_text(self.epilog)

                # determine help from format above
                return formatter.format_help()

class CoreArgumentParser(ArgumentParser):

        def format_help(self):
                formatter = self._get_formatter()

                # usage
                formatter.add_usage(self.usage, self._actions,
                                    self._mutually_exclusive_groups)

                # description
                formatter.add_text(self.description)

                #Add command help (not managed by argparse but by python cmd2)
                formatter.start_section("commands")
                formatter.add_text("Available commands")
                formatter.add_arguments(self.actions)
                formatter.end_section()

                # positionals, optionals and user-defined groups
                for action_group in self._action_groups:
                        if action_group.title=="positional arguments":
                                newGroupActions=[]
                                #search for the cmds positionals argument to remove it from the help message
                                for action in action_group._group_actions:
                                        if action.dest!="cmds":
                                               newGroupActions.append(action)
                                action_group._group_actions=newGroupActions
                        formatter.start_section(action_group.title)
                        formatter.add_text(action_group.description)
                        formatter.add_arguments(action_group._group_actions)
                        formatter.end_section()

                # epilog
                formatter.add_text(self.epilog)

                # determine help from format above
                return formatter.format_help()
   <commit_msg>Added parse_args wrapper<commit_after>__author__="UShareSoft"
__date__ ="$Jan 21, 2014 5:48:54 PM$"

import argparse

class ArgumentParserError(Exception): pass

class ArgumentParser(argparse.ArgumentParser):

	def error(self, message):
        	raise ArgumentParserError(message)

        def format_help(self):
                formatter = self._get_formatter()

                # description
                formatter.add_text(self.description)

                # usage
                formatter.add_usage(self.usage, self._actions,
                                    self._mutually_exclusive_groups)

                # positionals, optionals and user-defined groups
                for action_group in self._action_groups:
                    formatter.start_section(action_group.title)
                    formatter.add_text(action_group.description)
                    formatter.add_arguments(action_group._group_actions)
                    formatter.end_section()

                # epilog
                formatter.add_text(self.epilog)

                # determine help from format above
                return formatter.format_help()

	# Catch the system exit exception and return an empty object if the helper is called
	def parse_args(self, args):
		try:
			super(ArgumentParser,self).parse_args(args)
		except SystemExit as e :
                        return

class CoreArgumentParser(ArgumentParser):

        def format_help(self):
                formatter = self._get_formatter()

                # usage
                formatter.add_usage(self.usage, self._actions,
                                    self._mutually_exclusive_groups)

                # description
                formatter.add_text(self.description)

                #Add command help (not managed by argparse but by python cmd2)
                formatter.start_section("commands")
                formatter.add_text("Available commands")
                formatter.add_arguments(self.actions)
                formatter.end_section()

                # positionals, optionals and user-defined groups
                for action_group in self._action_groups:
                        if action_group.title=="positional arguments":
                                newGroupActions=[]
                                #search for the cmds positionals argument to remove it from the help message
                                for action in action_group._group_actions:
                                        if action.dest!="cmds":
                                               newGroupActions.append(action)
                                action_group._group_actions=newGroupActions
                        formatter.start_section(action_group.title)
                        formatter.add_text(action_group.description)
                        formatter.add_arguments(action_group._group_actions)
                        formatter.end_section()

                # epilog
                formatter.add_text(self.epilog)

                # determine help from format above
                return formatter.format_help()

<|endoftext|>


<commit_before>"""plot_overrepresented.py

Plot the number of over-represented blockgroups versus the number of households.
"""
from __future__ import division
import math
import csv
import numpy as np
from matplotlib import pylab as plt
import networkx as nx
from scipy.stats import linregress

colours = {'Lower':'#4F8F6B',
        'Higher':'#C1A62E',
        'Middle':'#4B453C'}

#
# Import data
#

## List of MSA
msa = {}
with open('data/names/msa.csv', 'r') as source:
    reader = csv.reader(source, delimiter='\t')
    reader.next()
    for rows in reader:
        msa[rows[0]] = rows[1]

## Classes
classes = {}
with open('extr/classes/msa_average/classes.csv', 'r') as source:
    reader = csv.reader(source, delimiter='\t')
    reader.next()
    for rows in reader:
        classes[rows[0]] =[int(r) for r in rows[1:]]

## Number of neighbourhoods
number_neigh = {cl:{} for cl in classes}
with open('extr/units/numbers/numbers.csv', 'r') as source:
    reader =  csv.reader(source, delimiter='\t')
    cats = reader.next()[1:]
    for rows in reader:
        for i,clust in enumerate(rows[1:]):
            number_neigh[cats[i]][rows[0]] = float(clust)

## Number of households per MSA
households = {}
with open('data/population/msa/msa_households.csv', 'r') as source:
    reader = csv.reader(source, delimiter='\t')
    reader.next()
    for rows in reader:
        households[rows[0]] = int(rows[1])

#
# Aggregate and fit
#
number = {cl:[number_neigh[cl][c] for c in households] for cl in classes}
population = [households[c] for c in households]

slopes = {}
r_values = {}
intercepts = {}
for cl in classes:
    print "Power-law fit for %s income class"%cl
    slope, intercept, r_value, p_value, std_err = linregress([math.log(p) for
        p in population],[math.log(d) for d in number[cl]])
    slopes[cl] = slope
    r_values[cl] = r_value
    intercepts[cl] = intercept
    print "alpha = %s (R^2=%s)"%(slope, r_value)

#
# Plot the distribution
#
fig = plt.figure(figsize=(24,8))
for i,cl in enumerate(classes):
    ax = fig.add_subplot(1, len(classes), i+1)
    ax.plot(population, number[cl], 'o', color=colours[cl], lw=2,
            mec=colours[cl],label=r"$%s$"%cl)
    ax.plot(sorted(population),
            [math.exp(intercepts[cl])*p**slopes[cl] for p in sorted(population)],
            linestyle = '--',
            color='black',
            label=r'$N_o \sim H^{\,%.2f}$'%(slopes[cl]))
    ax.set_xlabel(r'$H$', fontsize=20)
    ax.set_ylabel(r'$N_{o}$', fontsize=20)
    ax.set_xscale('log')
    ax.set_yscale('log')
    ax.spines['top'].set_visible(False)
    ax.spines['right'].set_visible(False)
    ax.spines['left'].set_position(('outward', 10))  # outward by 10 points
    ax.spines['bottom'].set_position(('outward', 10))  # outward by 10 points
    ax.spines['left'].set_smart_bounds(True)
    ax.spines['bottom'].set_smart_bounds(True)
    ax.yaxis.set_ticks_position('left')
    ax.xaxis.set_ticks_position('bottom')
    ax.legend(loc='upper left', numpoints=1, frameon=False)
plt.savefig('figures/paper/number_overrepresented.pdf', bbox_inches='tight')
plt.show()

<commit_msg>Save figure to si figures<commit_after>"""plot_overrepresented.py

Plot the number of over-represented blockgroups versus the number of households.
"""
from __future__ import division
import math
import csv
import numpy as np
from matplotlib import pylab as plt
import networkx as nx
from scipy.stats import linregress

colours = {'Lower':'#4F8F6B',
        'Higher':'#C1A62E',
        'Middle':'#4B453C'}

#
# Import data
#

## List of MSA
msa = {}
with open('data/names/msa.csv', 'r') as source:
    reader = csv.reader(source, delimiter='\t')
    reader.next()
    for rows in reader:
        msa[rows[0]] = rows[1]

## Classes
classes = {}
with open('extr/classes/msa_average/classes.csv', 'r') as source:
    reader = csv.reader(source, delimiter='\t')
    reader.next()
    for rows in reader:
        classes[rows[0]] =[int(r) for r in rows[1:]]

## Number of neighbourhoods
number_neigh = {cl:{} for cl in classes}
with open('extr/units/numbers/numbers.csv', 'r') as source:
    reader =  csv.reader(source, delimiter='\t')
    cats = reader.next()[1:]
    for rows in reader:
        for i,clust in enumerate(rows[1:]):
            number_neigh[cats[i]][rows[0]] = float(clust)

## Number of households per MSA
households = {}
with open('data/population/msa/msa_households.csv', 'r') as source:
    reader = csv.reader(source, delimiter='\t')
    reader.next()
    for rows in reader:
        households[rows[0]] = int(rows[1])

#
# Aggregate and fit
#
number = {cl:[number_neigh[cl][c] for c in households] for cl in classes}
population = [households[c] for c in households]

slopes = {}
r_values = {}
intercepts = {}
for cl in classes:
    print "Power-law fit for %s income class"%cl
    slope, intercept, r_value, p_value, std_err = linregress([math.log(p) for
        p in population],[math.log(d) for d in number[cl]])
    slopes[cl] = slope
    r_values[cl] = r_value
    intercepts[cl] = intercept
    print "alpha = %s (R^2=%s)"%(slope, r_value)

#
# Plot the distribution
#
fig = plt.figure(figsize=(24,8))
for i,cl in enumerate(classes):
    ax = fig.add_subplot(1, len(classes), i+1)
    ax.plot(population, number[cl], 'o', color=colours[cl], lw=2,
            mec=colours[cl],label=r"$%s$"%cl)
    ax.plot(sorted(population),
            [math.exp(intercepts[cl])*p**slopes[cl] for p in sorted(population)],
            linestyle = '--',
            color='black',
            label=r'$N_o \sim H^{\,%.2f}$'%(slopes[cl]))
    ax.set_xlabel(r'$H$', fontsize=20)
    ax.set_ylabel(r'$N_{o}$', fontsize=20)
    ax.set_xscale('log')
    ax.set_yscale('log')
    ax.spines['top'].set_visible(False)
    ax.spines['right'].set_visible(False)
    ax.spines['left'].set_position(('outward', 10))  # outward by 10 points
    ax.spines['bottom'].set_position(('outward', 10))  # outward by 10 points
    ax.spines['left'].set_smart_bounds(True)
    ax.spines['bottom'].set_smart_bounds(True)
    ax.yaxis.set_ticks_position('left')
    ax.xaxis.set_ticks_position('bottom')
    ax.legend(loc='upper left', numpoints=1, frameon=False)
plt.savefig('figures/paper/si/number_overrepresented.pdf', bbox_inches='tight')
plt.show()

<|endoftext|>