<commit_before>#!/usr/bin/env python
# encoding: utf-8
"""
update_cases.py

Update cases with information from a old instance

Created by Måns Magnusson on 2016-05-11.
Copyright (c) 2016 ScoutTeam__. All rights reserved.

"""
import logging

import click
from mongoengine import DoesNotExist
import yaml

from scout.models import Variant, Event
from scout.models.phenotype_term import PhenotypeTerm

logger = logging.getLogger(__name__)

def update_variants(adapter, case_obj, old_variants):
    """Update new variants with old information."""
    for variant in old_variants:
        new_variant = Variant.objects(variant_id=variant['variant_id']).first()
        if new_variant is None:
            logger.warn("missing variant: %s", variant['variant_id'])
            continue
        if variant.get('manual_rank'):
            logger.info("updating manual rank: %s", variant['manual_rank'])
            new_variant.manual_rank = variant['manual_rank']
        if variant.get('sanger_ordered') is True:
            logger.info("marking sanger order")
            new_variant.sanger_ordered = True
            if new_variant not in case_obj.suspects:
                logger.info("adding suspect: %s", variant['variant_id'])
                new_variant.suspects.append(new_variant)
        yield new_variant

def update_events(adapter, case_obj, old_events):
    """Update old events (comments) for a new case."""
    for event in old_events:
        if not event.get('content'):
            logger.debug("skipping event with no content: %s", event['verb'])
            continue
        try:
            user_obj = adapter.user(event['author'])
        except DoesNotExist:
            logger.warn("unable to find user: %s", event['author'])
            continue

        institute = adapter.institute(case_obj.owner)
        new_event = Event(case=case_obj, institute=institute, link=event['link'],
                          category=event['category'], author=user_obj,
                          subject=event['subject'], verb=event['verb'],
                          level=event['level'], content=event['content'],
                          created_at=event['created_at'])
        if event['variant_id']:
            new_event.variant_id = event['variant_id']
        yield new_event

def update_case(adapter, case_obj, exported_data):
    """docstring for update_case"""
    case_id = case_obj.case_id

    # Update the collaborators
    if exported_data['collaborators']:
        collaborators = set([coll for coll in case_obj.collaborators] +
                            exported_data['collaborators'])
        logger.info("Set collaborators to %s ", ', '.join(collaborators))
        case_obj.collaborators = list(collaborators)

    # Set assignee
    if case_obj.assignee is None:
        if exported_data['assignee']:
            mail = exported_data['assignee']
            try:
                user_obj = adapter.user(email=mail)
                logger.info("Assigning user %s", mail)
                case_obj.assignee = user_obj
            except DoesNotExist:
                logger.info("Could not find user user %s", mail)

    # Add the old suspects
    if exported_data['suspects']:
        existing_suspects = case_obj.suspects
        for suspect in exported_data['suspects']:
            variant_obj = Variant.objects.get(variant_id=suspect, case_id=case_id)
            if variant_obj:
                logger.info("Adding suspect %s", variant_obj.display_name)
                existing_suspects.append(variant_obj)
            else:
                logger.info("Could not find suspect %s in database", suspect)
        case_obj.suspects = existing_suspects

    # Add the old suspects
    if exported_data['causatives']:
        existing_causatives = case_obj.causatives
        for causative in exported_data['causatives']:
            variant_obj = Variant.objects.get(variant_id=causative, case_id=case_id)
            if variant_obj:
                logger.info("Adding causative %s", variant_obj.display_name)
                existing_causatives.append(variant_obj)
            else:
                logger.info("Could not find causative %s in database", causative)
        case_obj.causatives = existing_causatives

    # Update the synopsis
    if exported_data['synopsis']:
        if not case_obj.synopsis:
            logger.info("Updating synopsis")
            case_obj.synopsis = exported_data['synopsis']

    # Update the phenotype terms
    # If the case is phenotyped, skip to add the terms
    phenotype_terms = exported_data['phenotype_terms']
    if phenotype_terms:
        logger.info("Updating phenotype terms")
        existing_terms = case_obj.phenotype_terms
        existing_ids = set(term.phenotype_id for term in existing_terms)
        for hpo_id in phenotype_terms:
            if hpo_id in existing_ids:
                logger.info("term already added: %s", hpo_id)
            else:
                hpo_obj = adapter.hpo_term(hpo_id)
                if hpo_obj:
                    new_term = PhenotypeTerm(phenotype_id=hpo_id,
                                             feature=hpo_obj.description)
                    logger.info("Adding term %s", hpo_id)
                    case_obj.phenotype_terms.append(new_term)
                else:
                    logger.info("Could not find term %s", hpo_id)

    # Update the phenotype groups
    # If the case has a phenotype group, skip to add the groups
    phenotype_groups = exported_data['phenotype_groups']
    if phenotype_groups:
        logger.info("Updating phenotype groups")
        existing_groups = case_obj.phenotype_groups
        if len(existing_groups) == 0:
            for term in phenotype_groups:
                hpo_obj = adapter.hpo_term(term)
                if hpo_obj:
                    new_term = PhenotypeTerm(phenotype_id=term,
                                             feature=hpo_obj.description)
                    logger.info("Adding term %s", term)
                    existing_groups.append(new_term)
                else:
                    logger.info("Could not find term %s", term)
            case_obj.phenotype_groups = existing_groups
        else:
            logger.info("Case already had phenotype groups")

    logger.info("Saving case")
    case_obj.save()

@click.command('update_cases', short_help='Update cases')
@click.argument('exported_cases', type=click.File('r'))
@click.pass_context
def update_cases(context, exported_cases):
    """Update all information that was manually annotated from a old instance

        Takes a yaml file
    """
    logger.info("Running scout update cases")
    exported_data = yaml.load(exported_cases)
    adapter = context.obj['adapter']
    for case_obj in adapter.cases():
        logger.debug("working on case: %s", case_obj.case_id)
        if not case_obj.is_migrated:
            case_id = case_obj.case_id
            if case_id in exported_data:
                logger.info("Updating case: %s" % case_id)
                exported_info = exported_data[case_id]
                update_case(adapter, case_obj, exported_info)
                if exported_info.get('variants'):
                    logger.debug("updating variants")
                    new_variants = update_variants(adapter, case_obj,
                                                   exported_info['variants'])
                    for new_variant in new_variants:
                        new_variant.save()
                if exported_info.get('events'):
                    logger.debug("updating events")
                    new_events = update_events(adapter, case_obj,
                                               exported_info['events'])
                    for new_event in new_events:
                        logger.info("adding new event: %s", new_event.verb)
                        new_event.save()
            case_obj.is_migrated = True
            case_obj.save()
<commit_msg>fix mark is_migrated<commit_after>#!/usr/bin/env python
# encoding: utf-8
"""
update_cases.py

Update cases with information from a old instance

Created by Måns Magnusson on 2016-05-11.
Copyright (c) 2016 ScoutTeam__. All rights reserved.

"""
import logging

import click
from mongoengine import DoesNotExist
import yaml

from scout.models import Variant, Event
from scout.models.phenotype_term import PhenotypeTerm

logger = logging.getLogger(__name__)

def update_variants(adapter, case_obj, old_variants):
    """Update new variants with old information."""
    for variant in old_variants:
        new_variant = Variant.objects(variant_id=variant['variant_id']).first()
        if new_variant is None:
            logger.warn("missing variant: %s", variant['variant_id'])
            continue
        if variant.get('manual_rank'):
            logger.info("updating manual rank: %s", variant['manual_rank'])
            new_variant.manual_rank = variant['manual_rank']
        if variant.get('sanger_ordered') is True:
            logger.info("marking sanger order")
            new_variant.sanger_ordered = True
            if new_variant not in case_obj.suspects:
                logger.info("adding suspect: %s", variant['variant_id'])
                new_variant.suspects.append(new_variant)
        yield new_variant

def update_events(adapter, case_obj, old_events):
    """Update old events (comments) for a new case."""
    for event in old_events:
        if not event.get('content'):
            logger.debug("skipping event with no content: %s", event['verb'])
            continue
        try:
            user_obj = adapter.user(event['author'])
        except DoesNotExist:
            logger.warn("unable to find user: %s", event['author'])
            continue

        institute = adapter.institute(case_obj.owner)
        new_event = Event(case=case_obj, institute=institute, link=event['link'],
                          category=event['category'], author=user_obj,
                          subject=event['subject'], verb=event['verb'],
                          level=event['level'], content=event['content'],
                          created_at=event['created_at'])
        if event['variant_id']:
            new_event.variant_id = event['variant_id']
        yield new_event

def update_case(adapter, case_obj, exported_data):
    """docstring for update_case"""
    case_id = case_obj.case_id

    # Update the collaborators
    if exported_data['collaborators']:
        collaborators = set([coll for coll in case_obj.collaborators] +
                            exported_data['collaborators'])
        logger.info("Set collaborators to %s ", ', '.join(collaborators))
        case_obj.collaborators = list(collaborators)

    # Set assignee
    if case_obj.assignee is None:
        if exported_data['assignee']:
            mail = exported_data['assignee']
            try:
                user_obj = adapter.user(email=mail)
                logger.info("Assigning user %s", mail)
                case_obj.assignee = user_obj
            except DoesNotExist:
                logger.info("Could not find user user %s", mail)

    # Add the old suspects
    if exported_data['suspects']:
        existing_suspects = case_obj.suspects
        for suspect in exported_data['suspects']:
            variant_obj = Variant.objects.get(variant_id=suspect, case_id=case_id)
            if variant_obj:
                logger.info("Adding suspect %s", variant_obj.display_name)
                existing_suspects.append(variant_obj)
            else:
                logger.info("Could not find suspect %s in database", suspect)
        case_obj.suspects = existing_suspects

    # Add the old suspects
    if exported_data['causatives']:
        existing_causatives = case_obj.causatives
        for causative in exported_data['causatives']:
            variant_obj = Variant.objects.get(variant_id=causative, case_id=case_id)
            if variant_obj:
                logger.info("Adding causative %s", variant_obj.display_name)
                existing_causatives.append(variant_obj)
            else:
                logger.info("Could not find causative %s in database", causative)
        case_obj.causatives = existing_causatives

    # Update the synopsis
    if exported_data['synopsis']:
        if not case_obj.synopsis:
            logger.info("Updating synopsis")
            case_obj.synopsis = exported_data['synopsis']

    # Update the phenotype terms
    # If the case is phenotyped, skip to add the terms
    phenotype_terms = exported_data['phenotype_terms']
    if phenotype_terms:
        logger.info("Updating phenotype terms")
        existing_terms = case_obj.phenotype_terms
        existing_ids = set(term.phenotype_id for term in existing_terms)
        for hpo_id in phenotype_terms:
            if hpo_id in existing_ids:
                logger.info("term already added: %s", hpo_id)
            else:
                hpo_obj = adapter.hpo_term(hpo_id)
                if hpo_obj:
                    new_term = PhenotypeTerm(phenotype_id=hpo_id,
                                             feature=hpo_obj.description)
                    logger.info("Adding term %s", hpo_id)
                    case_obj.phenotype_terms.append(new_term)
                else:
                    logger.info("Could not find term %s", hpo_id)

    # Update the phenotype groups
    # If the case has a phenotype group, skip to add the groups
    phenotype_groups = exported_data['phenotype_groups']
    if phenotype_groups:
        logger.info("Updating phenotype groups")
        existing_groups = case_obj.phenotype_groups
        if len(existing_groups) == 0:
            for term in phenotype_groups:
                hpo_obj = adapter.hpo_term(term)
                if hpo_obj:
                    new_term = PhenotypeTerm(phenotype_id=term,
                                             feature=hpo_obj.description)
                    logger.info("Adding term %s", term)
                    existing_groups.append(new_term)
                else:
                    logger.info("Could not find term %s", term)
            case_obj.phenotype_groups = existing_groups
        else:
            logger.info("Case already had phenotype groups")

    logger.info("Saving case")
    case_obj.save()

@click.command('update_cases', short_help='Update cases')
@click.argument('exported_cases', type=click.File('r'))
@click.pass_context
def update_cases(context, exported_cases):
    """Update all information that was manually annotated from a old instance

        Takes a yaml file
    """
    logger.info("Running scout update cases")
    exported_data = yaml.load(exported_cases)
    adapter = context.obj['adapter']
    for case_obj in adapter.cases():
        logger.debug("working on case: %s", case_obj.case_id)
        if not case_obj.is_migrated:
            case_id = case_obj.case_id
            if case_id in exported_data:
                logger.info("Updating case: %s" % case_id)
                exported_info = exported_data[case_id]
                update_case(adapter, case_obj, exported_info)
                if exported_info.get('variants'):
                    logger.debug("updating variants")
                    new_variants = update_variants(adapter, case_obj,
                                                   exported_info['variants'])
                    for new_variant in new_variants:
                        new_variant.save()
                if exported_info.get('events'):
                    logger.debug("updating events")
                    new_events = update_events(adapter, case_obj,
                                               exported_info['events'])
                    for new_event in new_events:
                        logger.info("adding new event: %s", new_event.verb)
                        new_event.save()

                case_obj.is_migrated = True
                case_obj.save()
<|endoftext|>


<commit_before>#!/usr/bin/env python
# -*- Mode: python; tab-width: 4; indent-tabs-mode: nil -*-
#
# This program is free software; you can redistribute it and/or modify it under
# the terms of the GNU Lesser General Public License as published by the Free
# Software Foundation; either version 2 of the License, or (at your option) any
# later version.
#
# This program is distributed in the hope that it will be useful, but WITHOUT
# ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS
# FOR A PARTICULAR PURPOSE.  See the GNU Lesser General Public License for more
# details.
#
# You should have received a copy of the GNU Lesser General Public License along
# with this program; if not, write to the Free Software Foundation, Inc., 51
# Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA.
#
# Copyright (C) 2012 Lanedo GmbH
#

import string

import utils
import VariableFactory
import TypeFactory

"""
The Field class takes care of handling Input and Output TLVs
"""
class Field:

    """
    Constructor
    """
    def __init__(self, prefix, dictionary, common_objects_dictionary):
        # The field prefix, usually the name of the Container,
        #  e.g. "Qmi Message Ctl Something Output"
        self.prefix = prefix
        # The name of the specific field, e.g. "Result"
        self.name = dictionary['name']
        # The specific TLV ID
        self.id = dictionary['id']
        # Whether the field is to be considered mandatory in the message
        self.mandatory = dictionary['mandatory']
        # The type, which must always be "TLV"
        self.type = dictionary['type']

        # Create the composed full name (prefix + name),
        #  e.g. "Qmi Message Ctl Something Output Result"
        self.fullname = dictionary['fullname'] if 'fullname' in dictionary else self.prefix + ' ' + self.name

        # Create our variable object
        self.variable = VariableFactory.create_variable(dictionary, self.fullname)

        # Create the variable name within the Container
        self.variable_name = 'arg_' + string.lower(utils.build_underscore_name(self.name))

        # Create the ID enumeration name
        self.id_enum_name = string.upper(utils.build_underscore_name(self.prefix + ' TLV ' + self.name))

        # Output Fields may have prerequisites
        self.prerequisites = []
        if 'prerequisites' in dictionary:
            self.prerequisites = dictionary['prerequisites']
            # First, look for references to common types
            for prerequisite_dictionary in self.prerequisites:
                if 'common-ref' in prerequisite_dictionary:
                    for common in common_objects_dictionary:
                        if common['type'] == 'prerequisite' and \
                           common['common-ref'] == prerequisite_dictionary['common-ref']:
                           # Replace the reference with a copy of the common dictionary
                           copy = dict(common)
                           self.prerequisites.remove(prerequisite_dictionary)
                           self.prerequisites.append(copy)
                           break
                    else:
                        raise RuntimeError('Common type \'%s\' not found' % prerequisite_dictionary['name'])

    """
    Emit new types required by this field
    """
    def emit_types(self, hfile, cfile):
        if TypeFactory.is_type_emitted(self.fullname) is False:
            TypeFactory.set_type_emitted(self.fullname)
            self.variable.emit_types(hfile)
            self.variable.emit_helper_methods(hfile, cfile)

    """
    Emit the method responsible for getting this TLV from the input/output
    container
    """
    def emit_getter(self, hfile, cfile):
        input_variable_name = utils.build_underscore_name(self.name)
        variable_getter_dec = self.variable.build_getter_declaration('    ', input_variable_name)
        variable_getter_doc = self.variable.build_getter_documentation(' * ', input_variable_name)
        variable_getter_imp = self.variable.build_getter_implementation('    ', 'self->' + self.variable_name, input_variable_name, True)
        translations = { 'name'                : self.name,
                         'variable_name'       : self.variable_name,
                         'variable_getter_dec' : variable_getter_dec,
                         'variable_getter_doc' : variable_getter_doc,
                         'variable_getter_imp' : variable_getter_imp,
                         'underscore'          : utils.build_underscore_name(self.name),
                         'prefix_camelcase'    : utils.build_camelcase_name(self.prefix),
                         'prefix_underscore'   : utils.build_underscore_name(self.prefix) }

        # Emit the getter header
        template = (
            '\n'
            'gboolean ${prefix_underscore}_get_${underscore} (\n'
            '    ${prefix_camelcase} *self,\n'
            '${variable_getter_dec}'
            '    GError **error);\n')
        hfile.write(string.Template(template).substitute(translations))

        # Emit the getter source
        template = (
            '\n'
            '/**\n'
            ' * ${prefix_underscore}_get_${underscore}:\n'
            ' * @self: a ${prefix_camelcase}.\n'
            '${variable_getter_doc}'
            ' * @error: a #GError.\n'
            ' *\n'
            ' * Get the \'${name}\' field from @self.\n'
            ' *\n'
            ' * Returns: #TRUE if the field is found, #FALSE otherwise.\n'
            ' */\n'
            'gboolean\n'
            '${prefix_underscore}_get_${underscore} (\n'
            '    ${prefix_camelcase} *self,\n'
            '${variable_getter_dec}'
            '    GError **error)\n'
            '{\n'
            '    g_return_val_if_fail (self != NULL, FALSE);\n'
            '\n'
            '    if (!self->${variable_name}_set) {\n'
            '        g_set_error (error,\n'
            '                     QMI_CORE_ERROR,\n'
            '                     QMI_CORE_ERROR_TLV_NOT_FOUND,\n'
            '                     "Field \'${name}\' was not found in the message");\n'
            '        return FALSE;\n'
            '    }\n'
            '\n'
            '${variable_getter_imp}'
            '\n'
            '    return TRUE;\n'
            '}\n')
        cfile.write(string.Template(template).substitute(translations))

    """
    Emit the method responsible for setting this TLV in the input/output
    container
    """
    def emit_setter(self, hfile, cfile):
        input_variable_name = utils.build_underscore_name(self.name)
        variable_setter_dec = self.variable.build_setter_declaration('    ', input_variable_name)
        variable_setter_doc = self.variable.build_setter_documentation(' * ', input_variable_name)
        variable_setter_imp = self.variable.build_setter_implementation('    ', input_variable_name, 'self->' + self.variable_name)
        translations = { 'name'                : self.name,
                         'variable_name'       : self.variable_name,
                         'variable_setter_dec' : variable_setter_dec,
                         'variable_setter_doc' : variable_setter_doc,
                         'variable_setter_imp' : variable_setter_imp,
                         'underscore'          : utils.build_underscore_name(self.name),
                         'prefix_camelcase'    : utils.build_camelcase_name(self.prefix),
                         'prefix_underscore'   : utils.build_underscore_name(self.prefix) }

        # Emit the setter header
        template = (
            '\n'
            'gboolean ${prefix_underscore}_set_${underscore} (\n'
            '    ${prefix_camelcase} *self,\n'
            '${variable_setter_dec}'
            '    GError **error);\n')
        hfile.write(string.Template(template).substitute(translations))

        # Emit the setter source
        template = (
            '\n'
            '/**\n'
            ' * ${prefix_underscore}_set_${underscore}:\n'
            ' * @self: a ${prefix_camelcase}.\n'
            '${variable_setter_doc}'
            ' * @error: a #GError.\n'
            ' *\n'
            ' * Set the \'${name}\' field in the message.\n'
            ' *\n'
            ' * Returns: #TRUE if @value was successfully set, #FALSE otherwise.\n'
            ' */\n'
            'gboolean\n'
            '${prefix_underscore}_set_${underscore} (\n'
            '    ${prefix_camelcase} *self,\n'
            '${variable_setter_dec}'
            '    GError **error)\n'
            '{\n'
            '    g_return_val_if_fail (self != NULL, FALSE);\n'
            '\n'
            '    self->${variable_name}_set = TRUE;\n'
            '${variable_setter_imp}'
            '\n'
            '    return TRUE;\n'
            '}\n')
        cfile.write(string.Template(template).substitute(translations))

    """
    Emit the code responsible for adding the TLV to the QMI message
    """
    def emit_input_tlv_add(self, f, line_prefix):
        translations = { 'name'          : self.name,
                         'tlv_id'        : self.id_enum_name,
                         'variable_name' : self.variable_name,
                         'lp'            : line_prefix }

        template = (
            '${lp}guint8 buffer[1024];\n'
            '${lp}guint16 buffer_len = 1024;\n'
            '${lp}guint8 *buffer_aux = buffer;\n'
            '\n')
        f.write(string.Template(template).substitute(translations))

        # Now, write the contents of the variable into the buffer
        self.variable.emit_buffer_write(f, line_prefix, 'input->' + self.variable_name, 'buffer_aux', 'buffer_len')

        template = (
            '\n'
            '${lp}if (!qmi_message_tlv_add (self,\n'
            '${lp}                          (guint8)${tlv_id},\n'
            '${lp}                          (1024 - buffer_len),\n'
            '${lp}                          buffer,\n'
            '${lp}                          error)) {\n'
            '${lp}    g_prefix_error (error, \"Couldn\'t set the ${name} TLV: \");\n'
            '${lp}    qmi_message_unref (self);\n'
            '${lp}    return NULL;\n'
            '${lp}}\n')
        f.write(string.Template(template).substitute(translations))

    """
    Emit the code responsible for checking prerequisites in output TLVs
    """
    def emit_output_prerequisite_check(self, f, line_prefix):
        if self.prerequisites == []:
            f.write('%s/* No Prerequisites for field */\n' % line_prefix)
            return

        for prerequisite in self.prerequisites:
            translations = { 'lp'                     : line_prefix,
                             'prerequisite_field'     : utils.build_underscore_name(prerequisite['field']),
                             'prerequisite_operation' : prerequisite['operation'],
                             'prerequisite_value'     : prerequisite['value'] }
            template = (
                '${lp}/* Prerequisite.... */\n'
                '${lp}if (!(self->arg_${prerequisite_field} ${prerequisite_operation} ${prerequisite_value}))\n'
                '${lp}    break;\n')
            f.write(string.Template(template).substitute(translations))

    """
    Emit the code responsible for retrieving the TLV from the QMI message
    """
    def emit_output_tlv_get(self, f, line_prefix):
        translations = { 'name'                 : self.name,
                         'container_underscore' : utils.build_underscore_name (self.prefix),
                         'tlv_id'               : self.id_enum_name,
                         'variable_name'        : self.variable_name,
                         'lp'                   : line_prefix,
                         'error'                : 'error' if self.mandatory == 'yes' else 'NULL'}

        template = (
            '${lp}guint8 *buffer;\n'
            '${lp}guint16 buffer_len;\n'
            '\n'
            '${lp}if (qmi_message_tlv_get (message,\n'
            '${lp}                         ${tlv_id},\n'
            '${lp}                         &buffer_len,\n'
            '${lp}                         &buffer,\n'
            '${lp}                         ${error})) {\n'
            '${lp}    self->${variable_name}_set = TRUE;\n'
            '\n')
        f.write(string.Template(template).substitute(translations))

        # Now, read the contents of the buffer into the variable
        self.variable.emit_buffer_read(f, line_prefix + '    ', 'self->' + self.variable_name, 'buffer', 'buffer_len')

        template = (
            '\n'
            '${lp}    /* The remaining size of the buffer needs to be 0 if we successfully read the TLV */\n'
            '${lp}    if (buffer_len > 0) {\n'
            '${lp}        g_warning ("Left \'%u\' bytes unread when getting the \'${name}\' TLV", buffer_len);\n'
            '${lp}    }\n')

        if self.mandatory == 'yes':
            template += (
                '${lp}} else {\n'
                '${lp}    g_prefix_error (error, \"Couldn\'t get the ${name} TLV: \");\n'
                '${lp}    ${container_underscore}_unref (self);\n'
                '${lp}    return NULL;\n'
                '${lp}}\n')
        else:
            template += (
                '${lp}}\n')
        f.write(string.Template(template).substitute(translations))

    """
    Emit the method responsible for creating a printable representation of the TLV
    """
    def emit_output_tlv_get_printable(self, f):
        if TypeFactory.is_get_printable_emitted(self.fullname):
            return

        TypeFactory.set_get_printable_emitted(self.fullname)

        translations = { 'name'       : self.name,
                         'tlv_id'     : self.id_enum_name,
                         'underscore' : utils.build_underscore_name (self.fullname) }
        template = (
            '\n'
            'static gchar *\n'
            '${underscore}_get_printable (\n'
            '    QmiMessage *message,\n'
            '    const gchar *line_prefix)\n'
            '{\n'
            '    guint8 *buffer;\n'
            '    guint16 buffer_len;\n'
            '\n'
            '    if (qmi_message_tlv_get (message,\n'
            '                             ${tlv_id},\n'
            '                             &buffer_len,\n'
            '                             &buffer,\n'
            '                             NULL)) {\n'
            '        GString *printable;\n'
            '\n'
            '        printable = g_string_new ("");\n')
        f.write(string.Template(template).substitute(translations))

        # Now, read the contents of the buffer into the printable representation
        self.variable.emit_get_printable(f, '        ', 'printable', 'buffer', 'buffer_len')

        template = (
            '\n'
            '        /* The remaining size of the buffer needs to be 0 if we successfully read the TLV */\n'
            '        if (buffer_len > 0) {\n'
            '            g_warning ("Left \'%u\' bytes unread when getting the \'${name}\' TLV as printable", buffer_len);\n'
            '        }\n'
            '\n'
            '        return g_string_free (printable, FALSE);\n'
            '    }\n'
            '\n'
            '    return NULL;\n'
            '}\n')
        f.write(string.Template(template).substitute(translations))
<commit_msg>qmi-codegen: mark variable as set only after really having set<commit_after>#!/usr/bin/env python
# -*- Mode: python; tab-width: 4; indent-tabs-mode: nil -*-
#
# This program is free software; you can redistribute it and/or modify it under
# the terms of the GNU Lesser General Public License as published by the Free
# Software Foundation; either version 2 of the License, or (at your option) any
# later version.
#
# This program is distributed in the hope that it will be useful, but WITHOUT
# ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS
# FOR A PARTICULAR PURPOSE.  See the GNU Lesser General Public License for more
# details.
#
# You should have received a copy of the GNU Lesser General Public License along
# with this program; if not, write to the Free Software Foundation, Inc., 51
# Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA.
#
# Copyright (C) 2012 Lanedo GmbH
#

import string

import utils
import VariableFactory
import TypeFactory

"""
The Field class takes care of handling Input and Output TLVs
"""
class Field:

    """
    Constructor
    """
    def __init__(self, prefix, dictionary, common_objects_dictionary):
        # The field prefix, usually the name of the Container,
        #  e.g. "Qmi Message Ctl Something Output"
        self.prefix = prefix
        # The name of the specific field, e.g. "Result"
        self.name = dictionary['name']
        # The specific TLV ID
        self.id = dictionary['id']
        # Whether the field is to be considered mandatory in the message
        self.mandatory = dictionary['mandatory']
        # The type, which must always be "TLV"
        self.type = dictionary['type']

        # Create the composed full name (prefix + name),
        #  e.g. "Qmi Message Ctl Something Output Result"
        self.fullname = dictionary['fullname'] if 'fullname' in dictionary else self.prefix + ' ' + self.name

        # Create our variable object
        self.variable = VariableFactory.create_variable(dictionary, self.fullname)

        # Create the variable name within the Container
        self.variable_name = 'arg_' + string.lower(utils.build_underscore_name(self.name))

        # Create the ID enumeration name
        self.id_enum_name = string.upper(utils.build_underscore_name(self.prefix + ' TLV ' + self.name))

        # Output Fields may have prerequisites
        self.prerequisites = []
        if 'prerequisites' in dictionary:
            self.prerequisites = dictionary['prerequisites']
            # First, look for references to common types
            for prerequisite_dictionary in self.prerequisites:
                if 'common-ref' in prerequisite_dictionary:
                    for common in common_objects_dictionary:
                        if common['type'] == 'prerequisite' and \
                           common['common-ref'] == prerequisite_dictionary['common-ref']:
                           # Replace the reference with a copy of the common dictionary
                           copy = dict(common)
                           self.prerequisites.remove(prerequisite_dictionary)
                           self.prerequisites.append(copy)
                           break
                    else:
                        raise RuntimeError('Common type \'%s\' not found' % prerequisite_dictionary['name'])

    """
    Emit new types required by this field
    """
    def emit_types(self, hfile, cfile):
        if TypeFactory.is_type_emitted(self.fullname) is False:
            TypeFactory.set_type_emitted(self.fullname)
            self.variable.emit_types(hfile)
            self.variable.emit_helper_methods(hfile, cfile)

    """
    Emit the method responsible for getting this TLV from the input/output
    container
    """
    def emit_getter(self, hfile, cfile):
        input_variable_name = utils.build_underscore_name(self.name)
        variable_getter_dec = self.variable.build_getter_declaration('    ', input_variable_name)
        variable_getter_doc = self.variable.build_getter_documentation(' * ', input_variable_name)
        variable_getter_imp = self.variable.build_getter_implementation('    ', 'self->' + self.variable_name, input_variable_name, True)
        translations = { 'name'                : self.name,
                         'variable_name'       : self.variable_name,
                         'variable_getter_dec' : variable_getter_dec,
                         'variable_getter_doc' : variable_getter_doc,
                         'variable_getter_imp' : variable_getter_imp,
                         'underscore'          : utils.build_underscore_name(self.name),
                         'prefix_camelcase'    : utils.build_camelcase_name(self.prefix),
                         'prefix_underscore'   : utils.build_underscore_name(self.prefix) }

        # Emit the getter header
        template = (
            '\n'
            'gboolean ${prefix_underscore}_get_${underscore} (\n'
            '    ${prefix_camelcase} *self,\n'
            '${variable_getter_dec}'
            '    GError **error);\n')
        hfile.write(string.Template(template).substitute(translations))

        # Emit the getter source
        template = (
            '\n'
            '/**\n'
            ' * ${prefix_underscore}_get_${underscore}:\n'
            ' * @self: a ${prefix_camelcase}.\n'
            '${variable_getter_doc}'
            ' * @error: a #GError.\n'
            ' *\n'
            ' * Get the \'${name}\' field from @self.\n'
            ' *\n'
            ' * Returns: #TRUE if the field is found, #FALSE otherwise.\n'
            ' */\n'
            'gboolean\n'
            '${prefix_underscore}_get_${underscore} (\n'
            '    ${prefix_camelcase} *self,\n'
            '${variable_getter_dec}'
            '    GError **error)\n'
            '{\n'
            '    g_return_val_if_fail (self != NULL, FALSE);\n'
            '\n'
            '    if (!self->${variable_name}_set) {\n'
            '        g_set_error (error,\n'
            '                     QMI_CORE_ERROR,\n'
            '                     QMI_CORE_ERROR_TLV_NOT_FOUND,\n'
            '                     "Field \'${name}\' was not found in the message");\n'
            '        return FALSE;\n'
            '    }\n'
            '\n'
            '${variable_getter_imp}'
            '\n'
            '    return TRUE;\n'
            '}\n')
        cfile.write(string.Template(template).substitute(translations))

    """
    Emit the method responsible for setting this TLV in the input/output
    container
    """
    def emit_setter(self, hfile, cfile):
        input_variable_name = utils.build_underscore_name(self.name)
        variable_setter_dec = self.variable.build_setter_declaration('    ', input_variable_name)
        variable_setter_doc = self.variable.build_setter_documentation(' * ', input_variable_name)
        variable_setter_imp = self.variable.build_setter_implementation('    ', input_variable_name, 'self->' + self.variable_name)
        translations = { 'name'                : self.name,
                         'variable_name'       : self.variable_name,
                         'variable_setter_dec' : variable_setter_dec,
                         'variable_setter_doc' : variable_setter_doc,
                         'variable_setter_imp' : variable_setter_imp,
                         'underscore'          : utils.build_underscore_name(self.name),
                         'prefix_camelcase'    : utils.build_camelcase_name(self.prefix),
                         'prefix_underscore'   : utils.build_underscore_name(self.prefix) }

        # Emit the setter header
        template = (
            '\n'
            'gboolean ${prefix_underscore}_set_${underscore} (\n'
            '    ${prefix_camelcase} *self,\n'
            '${variable_setter_dec}'
            '    GError **error);\n')
        hfile.write(string.Template(template).substitute(translations))

        # Emit the setter source
        template = (
            '\n'
            '/**\n'
            ' * ${prefix_underscore}_set_${underscore}:\n'
            ' * @self: a ${prefix_camelcase}.\n'
            '${variable_setter_doc}'
            ' * @error: a #GError.\n'
            ' *\n'
            ' * Set the \'${name}\' field in the message.\n'
            ' *\n'
            ' * Returns: #TRUE if @value was successfully set, #FALSE otherwise.\n'
            ' */\n'
            'gboolean\n'
            '${prefix_underscore}_set_${underscore} (\n'
            '    ${prefix_camelcase} *self,\n'
            '${variable_setter_dec}'
            '    GError **error)\n'
            '{\n'
            '    g_return_val_if_fail (self != NULL, FALSE);\n'
            '\n'
            '${variable_setter_imp}'
            '    self->${variable_name}_set = TRUE;\n'
            '\n'
            '    return TRUE;\n'
            '}\n')
        cfile.write(string.Template(template).substitute(translations))

    """
    Emit the code responsible for adding the TLV to the QMI message
    """
    def emit_input_tlv_add(self, f, line_prefix):
        translations = { 'name'          : self.name,
                         'tlv_id'        : self.id_enum_name,
                         'variable_name' : self.variable_name,
                         'lp'            : line_prefix }

        template = (
            '${lp}guint8 buffer[1024];\n'
            '${lp}guint16 buffer_len = 1024;\n'
            '${lp}guint8 *buffer_aux = buffer;\n'
            '\n')
        f.write(string.Template(template).substitute(translations))

        # Now, write the contents of the variable into the buffer
        self.variable.emit_buffer_write(f, line_prefix, 'input->' + self.variable_name, 'buffer_aux', 'buffer_len')

        template = (
            '\n'
            '${lp}if (!qmi_message_tlv_add (self,\n'
            '${lp}                          (guint8)${tlv_id},\n'
            '${lp}                          (1024 - buffer_len),\n'
            '${lp}                          buffer,\n'
            '${lp}                          error)) {\n'
            '${lp}    g_prefix_error (error, \"Couldn\'t set the ${name} TLV: \");\n'
            '${lp}    qmi_message_unref (self);\n'
            '${lp}    return NULL;\n'
            '${lp}}\n')
        f.write(string.Template(template).substitute(translations))

    """
    Emit the code responsible for checking prerequisites in output TLVs
    """
    def emit_output_prerequisite_check(self, f, line_prefix):
        if self.prerequisites == []:
            f.write('%s/* No Prerequisites for field */\n' % line_prefix)
            return

        for prerequisite in self.prerequisites:
            translations = { 'lp'                     : line_prefix,
                             'prerequisite_field'     : utils.build_underscore_name(prerequisite['field']),
                             'prerequisite_operation' : prerequisite['operation'],
                             'prerequisite_value'     : prerequisite['value'] }
            template = (
                '${lp}/* Prerequisite.... */\n'
                '${lp}if (!(self->arg_${prerequisite_field} ${prerequisite_operation} ${prerequisite_value}))\n'
                '${lp}    break;\n')
            f.write(string.Template(template).substitute(translations))

    """
    Emit the code responsible for retrieving the TLV from the QMI message
    """
    def emit_output_tlv_get(self, f, line_prefix):
        translations = { 'name'                 : self.name,
                         'container_underscore' : utils.build_underscore_name (self.prefix),
                         'tlv_id'               : self.id_enum_name,
                         'variable_name'        : self.variable_name,
                         'lp'                   : line_prefix,
                         'error'                : 'error' if self.mandatory == 'yes' else 'NULL'}

        template = (
            '${lp}guint8 *buffer;\n'
            '${lp}guint16 buffer_len;\n'
            '\n'
            '${lp}if (qmi_message_tlv_get (message,\n'
            '${lp}                         ${tlv_id},\n'
            '${lp}                         &buffer_len,\n'
            '${lp}                         &buffer,\n'
            '${lp}                         ${error})) {\n'
            '${lp}    self->${variable_name}_set = TRUE;\n'
            '\n')
        f.write(string.Template(template).substitute(translations))

        # Now, read the contents of the buffer into the variable
        self.variable.emit_buffer_read(f, line_prefix + '    ', 'self->' + self.variable_name, 'buffer', 'buffer_len')

        template = (
            '\n'
            '${lp}    /* The remaining size of the buffer needs to be 0 if we successfully read the TLV */\n'
            '${lp}    if (buffer_len > 0) {\n'
            '${lp}        g_warning ("Left \'%u\' bytes unread when getting the \'${name}\' TLV", buffer_len);\n'
            '${lp}    }\n')

        if self.mandatory == 'yes':
            template += (
                '${lp}} else {\n'
                '${lp}    g_prefix_error (error, \"Couldn\'t get the ${name} TLV: \");\n'
                '${lp}    ${container_underscore}_unref (self);\n'
                '${lp}    return NULL;\n'
                '${lp}}\n')
        else:
            template += (
                '${lp}}\n')
        f.write(string.Template(template).substitute(translations))

    """
    Emit the method responsible for creating a printable representation of the TLV
    """
    def emit_output_tlv_get_printable(self, f):
        if TypeFactory.is_get_printable_emitted(self.fullname):
            return

        TypeFactory.set_get_printable_emitted(self.fullname)

        translations = { 'name'       : self.name,
                         'tlv_id'     : self.id_enum_name,
                         'underscore' : utils.build_underscore_name (self.fullname) }
        template = (
            '\n'
            'static gchar *\n'
            '${underscore}_get_printable (\n'
            '    QmiMessage *message,\n'
            '    const gchar *line_prefix)\n'
            '{\n'
            '    guint8 *buffer;\n'
            '    guint16 buffer_len;\n'
            '\n'
            '    if (qmi_message_tlv_get (message,\n'
            '                             ${tlv_id},\n'
            '                             &buffer_len,\n'
            '                             &buffer,\n'
            '                             NULL)) {\n'
            '        GString *printable;\n'
            '\n'
            '        printable = g_string_new ("");\n')
        f.write(string.Template(template).substitute(translations))

        # Now, read the contents of the buffer into the printable representation
        self.variable.emit_get_printable(f, '        ', 'printable', 'buffer', 'buffer_len')

        template = (
            '\n'
            '        /* The remaining size of the buffer needs to be 0 if we successfully read the TLV */\n'
            '        if (buffer_len > 0) {\n'
            '            g_warning ("Left \'%u\' bytes unread when getting the \'${name}\' TLV as printable", buffer_len);\n'
            '        }\n'
            '\n'
            '        return g_string_free (printable, FALSE);\n'
            '    }\n'
            '\n'
            '    return NULL;\n'
            '}\n')
        f.write(string.Template(template).substitute(translations))
<|endoftext|>


<commit_before>"""
Linter for verifying werewolf character files

Checks for a number of problems that are specific to werewolf characters. The only
public entry point is the lint function.
"""

import re
from . import nwod

def lint(character, fix=False, *, strict=False, prefs=None):
    """
    Verify the more complex elements in a werewolf sheet.

    Checks for werewolf-specific problems within the rules blocks of the
    character sheet. The problems it checks for are as follows:

    1. Virtue and Vice are present (strict)
    2. Tribe is recognized
    3. Auspice is present if tribe is non-pure
    4. Auspice is not present if tribe is pure
    5. Auspice is recognized, if present

    Args:
        character (dict): Character data to lint
        fix (bool): Whether to automatically correct certain problems
        strict (bool): Whether to report non-critical errors and omissions
        prefs (Settings): Object storing the available auspices, tribes, and
            pure tribes.

    Returns:
        List of problem descriptions. If no problems were found, the list will
        be empty.
    """
    problems = []
    dirty = False

    # Ensure that tribe and auspice are correct
    problems.extend(lint_tribe(character, prefs))

    # If the character has no sheet, we cannot proceed
    if not character.has_path:
        problems.append('Missing path')
        return problems

    # Load the sheet for deep linting
    with open(character['path'], 'r') as char_file:
        data = char_file.read()

    # Check that they have a vice and a virtue
    if strict:
        problems.extend(nwod.lint_vice_virtue(data))

    if dirty and data:
        with open(character['path'], 'w') as char_file:
            char_file.write(data)

    return problems

def lint_tribe(character, prefs):
    """
    Lint tribe and auspice issues
    """
    problems = []
    all_tribes = prefs.get('werewolf.tribes') + prefs.get('werewolf.pure')

    # Make sure tribe is present and recognized
    tribe = character.get_first('tribe')
    if not tribe:
        return problems
    if not tribe.lower() in all_tribes:
        problems.append("Unrecognized tribe '{}'".format(tribe))

    # Get auspice
    auspice = character.get_first('auspice')

    if auspice:
        # Pure should not have an auspice
        if tribe in prefs.get('werewolf.pure'):
            problems.append('Auspice present, but werewolf is pure')
        # regardless, auspice must be recognized
        if auspice.lower() not in prefs.get('werewolf.auspices'):
            problems.append("Unrecognized auspice '{}'".format(auspice))
    else:
        # Non-pure must have an auspice
        if tribe in prefs.get('werewolf.tribes'):
            problems.append('Missing auspice')

    return problems
<commit_msg>Store tribe key for reuse<commit_after>"""
Linter for verifying werewolf character files

Checks for a number of problems that are specific to werewolf characters. The only
public entry point is the lint function.
"""

import re
from . import nwod

def lint(character, fix=False, *, strict=False, prefs=None):
    """
    Verify the more complex elements in a werewolf sheet.

    Checks for werewolf-specific problems within the rules blocks of the
    character sheet. The problems it checks for are as follows:

    1. Virtue and Vice are present (strict)
    2. Tribe is recognized
    3. Auspice is present if tribe is non-pure
    4. Auspice is not present if tribe is pure
    5. Auspice is recognized, if present

    Args:
        character (dict): Character data to lint
        fix (bool): Whether to automatically correct certain problems
        strict (bool): Whether to report non-critical errors and omissions
        prefs (Settings): Object storing the available auspices, tribes, and
            pure tribes.

    Returns:
        List of problem descriptions. If no problems were found, the list will
        be empty.
    """
    problems = []
    dirty = False

    # Ensure that tribe and auspice are correct
    problems.extend(lint_tribe(character, prefs))

    # If the character has no sheet, we cannot proceed
    if not character.has_path:
        problems.append('Missing path')
        return problems

    # Load the sheet for deep linting
    with open(character['path'], 'r') as char_file:
        data = char_file.read()

    # Check that they have a vice and a virtue
    if strict:
        problems.extend(nwod.lint_vice_virtue(data))

    if dirty and data:
        with open(character['path'], 'w') as char_file:
            char_file.write(data)

    return problems

def lint_tribe(character, prefs):
    """
    Lint tribe and auspice issues
    """
    problems = []
    all_tribes = prefs.get('werewolf.tribes') + prefs.get('werewolf.pure')

    # Make sure tribe is present and recognized
    tribe = character.get_first('tribe')
    if not tribe:
        return problems

    tribe_key = tribe.lower()
    if not tribe_key in all_tribes:
        problems.append("Unrecognized tribe '{}'".format(tribe))

    # Get auspice
    auspice = character.get_first('auspice')

    if auspice:
        # Pure should not have an auspice
        if tribe_key in prefs.get('werewolf.pure'):
            problems.append('Auspice present, but werewolf is Pure')
        # regardless, auspice must be recognized
        if auspice.lower() not in prefs.get('werewolf.auspices'):
            problems.append("Unrecognized auspice '{}'".format(auspice))
    else:
        # Non-pure must have an auspice
        if tribe_key in prefs.get('werewolf.tribes'):
            problems.append('Missing auspice')

    return problems
<|endoftext|>


<commit_before>from RepositoryClient import *
from sandbox_scripts.QualiEnvironmentUtils.Sandbox import *
import base64

try:
    import gitlab
    imported_gitlab = True
except Exception as e:
    imported_gitlab = False

class GitLabClient(RepositoryClient):
    # ----------------------------------
    # ----------------------------------
    def __init__(self, sandbox,repository_resource ):
        super(GitLabClient,self).__init__(sandbox, repository_resource)
        if imported_gitlab == False:
            self.sandbox.report_error('gitlab lib was not found', write_to_output_window=True)
        self.url = repository_resource.get_attribute("GitLab URL")
        self.token = repository_resource.get_attribute("GitLab Token")
        self.project_name = repository_resource.get_attribute("GitLab Project Name")
        self.repository_path = repository_resource.get_attribute("Repository Path")

    # ----------------------------------
    # ----------------------------------
    def download(self, source, destination):

        gl = gitlab.Gitlab(self.url, self.token)
        projid = 0
        projects = ''
        try:
            for project in gl.getall(gl.getprojects):
                #print project['name']
                if project['name'] == self.project_name:
                    projid = project['id']
                    break
                else:
                    projects += str(project['id']) + '-' + project['name'] + '\n'
        except Exception as ex:
            raise QualiError("GitLabClient", "ERROR: Could not access repository at %s" % self.url + " : " + str(ex.message))

        if projid == 0:
            raise QualiError("GitLabClient","ERROR: Failed to locate project by name among \n" + projects)

        try:
            source = self.repository_path + source
            filebase64 = gl.getfile(projid, source, 'master')
            filetext = base64.b64decode(filebase64['content']).decode()
        except Exception as ex:
            raise QualiError("GitLabClient", "ERROR: Failed to retrieve file, which may be expected for " +
                             source + " : " + str(ex.message))

        try:
            with open(destination,'w') as dest:
                    dest.write(filetext)
            print "Downloaded: " + source
            return 0, "Successfully retrieved file from repository and saved to destination"
        except Exception as ex:
            raise QualiError("GitLabClient","ERROR: Retrieved file from repository - failed to save to destination : " + str(ex.message))
<commit_msg>import change, again<commit_after>from RepositoryClient import *
from sandbox_scripts.QualiEnvironmentUtils.Sandbox import *
import base64
import pip

try:
    import gitlab
    imported_gitlab = True
except:
    try:
        pip.main(["install","python-gitlab"])
        import gitlab
    except:
        imported_gitlab = False

class GitLabClient(RepositoryClient):
    # ----------------------------------
    # ----------------------------------
    def __init__(self, sandbox,repository_resource ):
        super(GitLabClient,self).__init__(sandbox, repository_resource)
        if imported_gitlab == False:
            self.sandbox.report_error('gitlab lib was not found', write_to_output_window=True)
        self.url = repository_resource.get_attribute("GitLab URL")
        self.token = repository_resource.get_attribute("GitLab Token")
        self.project_name = repository_resource.get_attribute("GitLab Project Name")
        self.repository_path = repository_resource.get_attribute("Repository Path")

    # ----------------------------------
    # ----------------------------------
    def download(self, source, destination):

        gl = gitlab.Gitlab(self.url, self.token)
        projid = 0
        projects = ''
        try:
            for project in gl.getall(gl.getprojects):
                #print project['name']
                if project['name'] == self.project_name:
                    projid = project['id']
                    break
                else:
                    projects += str(project['id']) + '-' + project['name'] + '\n'
        except Exception as ex:
            raise QualiError("GitLabClient", "ERROR: Could not access repository at %s" % self.url + " : " + str(ex.message))

        if projid == 0:
            raise QualiError("GitLabClient","ERROR: Failed to locate project by name among \n" + projects)

        try:
            source = self.repository_path + source
            filebase64 = gl.getfile(projid, source, 'master')
            filetext = base64.b64decode(filebase64['content']).decode()
        except Exception as ex:
            raise QualiError("GitLabClient", "ERROR: Failed to retrieve file, which may be expected for " +
                             source + " : " + str(ex.message))

        try:
            with open(destination,'w') as dest:
                    dest.write(filetext)
            print "Downloaded: " + source
            return 0, "Successfully retrieved file from repository and saved to destination"
        except Exception as ex:
            raise QualiError("GitLabClient","ERROR: Retrieved file from repository - failed to save to destination : " + str(ex.message))
<|endoftext|>


<commit_before>from collections import defaultdict

class StopWaitDispatch(Exception): pass
class StaticValue(object):
    def __init__(self, value):
        self.value = value

class Waiter(object):
    @property
    def wait_id(self):
        return hash(self)

    def process_fire(self, given):
        return StaticValue(given)

class StringWaiter(str, Waiter):
    def wait_id(self):
        return str(self)

class WaitPool(object):
    '''A structure that manages all `wait`ers, makes sure fired events
    get to the right places, and that all other waits are canceled when
    a one event is passed back to a generator.
    '''
    def __init__(self):
        self.waits = defaultdict(set)
        self.loop_refs = defaultdict(set)

    def wait(self, who, what):
        if type(what) is str:
            what = StringWaiter(what)

        self.waits[what.wait_id].add(who)
        self.loop_refs[who].add(what)
        return what.wait_id

    def fire(self, what, value):
        if type(what) is str:
            what = StringWaiter(what)

        static = False
        for handler in self.waits[what.wait_id]:
            if not static:
                try:
                    value = what.process_fire(value)
                except StopWaitDispatch:
                    break
                if type(value) == StaticValue:
                    static = True
                    value = value.value
            handler.fire_in(what.wait_id, value)

    def clear(self, who):
        for what in self.loop_refs[who]:
            self.waits[what.wait_id].remove(who)
        del self.loop_refs[who]
<commit_msg>Change event waiting code to account for new wait-type wrapping<commit_after>from collections import defaultdict

class StopWaitDispatch(Exception): pass
class StaticValue(object):
    def __init__(self, value):
        self.value = value

class Waiter(object):
    @property
    def wait_id(self):
        return hash(self)

    def process_fire(self, given):
        return StaticValue(given)

class StringWaiter(str, Waiter):
    def wait_id(self):
        return str(self)

class WaitPool(object):
    '''A structure that manages all `wait`ers, makes sure fired events
    get to the right places, and that all other waits are canceled when
    a one event is passed back to a generator.
    '''
    def __init__(self):
        self.waits = defaultdict(set)
        self.loop_refs = defaultdict(set)

    def wait(self, who, what):
        if isinstance(what, basestring):
            what = StringWaiter(what)

        self.waits[what.wait_id].add(who)
        self.loop_refs[who].add(what)
        return what.wait_id

    def fire(self, what, value):
        if isinstance(what, basestring):
            what = StringWaiter(what)

        static = False
        for handler in self.waits[what.wait_id]:
            if not static:
                try:
                    value = what.process_fire(value)
                except StopWaitDispatch:
                    break
                if type(value) == StaticValue:
                    static = True
                    value = value.value
            handler.fire_in(what.wait_id, value)

    def clear(self, who):
        for what in self.loop_refs[who]:
            self.waits[what.wait_id].remove(who)
        del self.loop_refs[who]
<|endoftext|>


<commit_before># Copyright (c) 2014 - 2017, Carlos Millett
# All rights reserved.
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions are met:
#
# 1. Redistributions of source code must retain the above copyright notice,
# this list of conditions and the following disclaimer.
#
# 2. Redistributions in binary form must reproduce the above copyright notice,
# this list of conditions and the following disclaimer in the documentation
# and/or other materials provided with the distribution.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
# ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE
# LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
# CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
# SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
# INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
# CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
# ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
# POSSIBILITY OF SUCH DAMAGE.

"""
Renamer renames your TV files into a nice new format.
The new filename includes the TV show name, season and episode numbers and
the episode title.
"""

#
# Import
#
from renamer import localpath
from renamer import web

# Information
__version__ = "0.2.0"
__author__ = "Carlos Millett"
__email__ = "carlos4735@gmail.com"
__maintainer__ = "Carlos Millett"
<commit_msg>version bump<commit_after># Copyright (c) 2014 - 2017, Carlos Millett
# All rights reserved.
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions are met:
#
# 1. Redistributions of source code must retain the above copyright notice,
# this list of conditions and the following disclaimer.
#
# 2. Redistributions in binary form must reproduce the above copyright notice,
# this list of conditions and the following disclaimer in the documentation
# and/or other materials provided with the distribution.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
# ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE
# LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
# CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
# SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
# INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
# CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
# ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
# POSSIBILITY OF SUCH DAMAGE.

"""
Renamer renames your TV files into a nice new format.
The new filename includes the TV show name, season and episode numbers and
the episode title.
"""

#
# Import
#
from renamer import localpath
from renamer import web

# Information
__version__ = "0.2.1"
__author__ = "Carlos Millett"
__email__ = "carlos4735@gmail.com"
__maintainer__ = "Carlos Millett"
<|endoftext|>


<commit_before>#!/usr/bin/env python
"""Tests for the autoload-credentials command."""

from __future__ import print_function

import argparse
import logging
import os
import pexpect
import sys

from textwrap import dedent

from jujupy import EnvJujuClient, JujuData
from utility import (
    configure_logging,
    enforce_juju_path,
    temp_dir,
    ensure_dir,
)

__metaclass__ = type

log = logging.getLogger("assess_autoload_credentials")

def assess_autoload_credentials(juju_bin):
    test_scenarios = [
        ('AWS using environment variables', aws_envvar_test_details),
        ('AWS using credentials file', aws_directory_test_details),
    ]

    for scenario_name, scenario_setup in test_scenarios:
        log.info('* Starting test scenario: {}'.format(scenario_name))
        test_autoload_credentials_stores_details(juju_bin, scenario_setup)

def test_autoload_credentials_stores_details(juju_bin, cloud_details_fn):
    """Test covering loading and storing credentials using autoload-credentials

    :param juju_bin: The full path to the juju binary to use for the test run.
    :param cloud_details_fn: A callable that takes the argument 'user' and
      'tmp_dir' that returns a tuple of:
        (dict -> environment variable changse,
        dict -> expected credential details)
      used to setup creation of credential details & comparison of the result.

    """
    user = 'testing_user'
    with temp_dir() as tmp_dir:
        client = EnvJujuClient.by_version(
            JujuData('local', juju_home=tmp_dir), juju_bin, False)

        env_var_changes, expected_details = cloud_details_fn(user, tmp_dir)
        # Inject well known username.
        env_var_changes.update({'USER': user})

        run_autoload_credentials(client, env_var_changes)

        client.env.load_yaml()

        assert_credentials_contains_expected_results(
            client.env.credentials,
            expected_details)

def assert_credentials_contains_expected_results(credentials, expected):
    if credentials != expected:
        raise ValueError(
            'Actual credentials do not match expected credentials.\n'
            'Expected: {expected}\nGot: {got}\n'.format(
                expected=expected,
                got=credentials
            )
        )

def run_autoload_credentials(client, envvars):
    """Execute the command 'juju autoload-credentials'.

    Simple interaction, calls juju autoload-credentials selects the first
    option and then quits.

    :param client: EnvJujuClient from which juju will be called.
    :param envvars: Dictionary containing environment variables to be used
      during execution.
      Note. Must contain a value for QUESTION_CLOUD_NAME to match against.
      Note. Must contain a value for SAVE_CLOUD_NAME to save against.

    """
    # Get juju path from client as we need to use it interactively.
    process = client.expect(
        'autoload-credentials', extra_env=envvars, include_e=False)
    process.expect(
        '.*1. \w+ credential "{}" \(new\).*'.format(
            envvars['QUESTION_CLOUD_NAME']))
    process.sendline('1')

    process.expect(
        'Enter cloud to which the credential belongs, or Q to quit.*')
    process.sendline(envvars['SAVE_CLOUD_NAME'])
    process.expect(
        'Saved aws credential "{}" to cloud \w+'.format(
            envvars['QUESTION_CLOUD_NAME']))
    process.sendline('q')
    process.expect(pexpect.EOF)

    if process.isalive():
        print(str(process))
        raise AssertionError('juju process failed to terminate')

def aws_envvar_test_details(
        user, tmp_dir, access_key='access_key', secret_key='secret_key'):
    env_var_changes = get_aws_environment(user, access_key, secret_key)

    expected_details = get_aws_expected_details_dict(
        user, access_key, secret_key)

    return env_var_changes, expected_details

def aws_directory_test_details(
        user, tmp_dir, access_key='access_key', secret_key='secret_key'
):
    expected_details = get_aws_expected_details_dict(
        'default', access_key, secret_key)

    write_aws_config_file(tmp_dir, access_key, secret_key)

    env_var_changes = dict(
        HOME=tmp_dir,
        SAVE_CLOUD_NAME='aws',
        QUESTION_CLOUD_NAME='default'
    )

    return env_var_changes, expected_details

def get_aws_expected_details_dict(cloud_name, access_key, secret_key):
    # Build credentials yaml file-like datastructure.
    return {
        'credentials': {
            'aws': {
                cloud_name: {
                    'auth-type': 'access-key',
                    'access-key': access_key,
                    'secret-key': secret_key,
                }
            }
        }
    }

def get_aws_environment(user, access_key, secret_key):
    """Return a dictionary containing keys suitable for AWS env vars.

    """
    return dict(
        SAVE_CLOUD_NAME='aws',
        QUESTION_CLOUD_NAME=user,
        AWS_ACCESS_KEY_ID=access_key,
        AWS_SECRET_ACCESS_KEY=secret_key
    )

def write_aws_config_file(tmp_dir, access_key, secret_key):
    """Write aws credentials file to tmp_dir

    :return: String path of created credentials file.

    """
    config_dir = os.path.join(tmp_dir, '.aws')
    config_file = os.path.join(config_dir, 'credentials')
    ensure_dir(config_dir)

    config_contents = dedent("""\
    [default]
    aws_access_key_id={}
    aws_secret_access_key={}
    """.format(access_key, secret_key))

    with open(config_file, 'w') as f:
        f.write(config_contents)

    return config_file

def parse_args(argv):
    parser = argparse.ArgumentParser(description="Test autoload-credentials.")
    parser.add_argument(
        'juju_bin', action=enforce_juju_path,
        help='Full path to the Juju binary.')
    parser.add_argument(
        '--verbose', action='store_const',
        default=logging.INFO, const=logging.DEBUG,
        help='Verbose test harness output.')

    return parser.parse_args(argv)

def main(argv=None):
    args = parse_args(argv)
    configure_logging(args.verbose)

    assess_autoload_credentials(args.juju_bin)
    return 0

if __name__ == '__main__':
    sys.exit(main())
<commit_msg>Ensure process is terminated if needed.<commit_after>#!/usr/bin/env python
"""Tests for the autoload-credentials command."""

from __future__ import print_function

import argparse
import logging
import os
import pexpect
import sys

from textwrap import dedent

from jujupy import EnvJujuClient, JujuData
from utility import (
    configure_logging,
    enforce_juju_path,
    temp_dir,
    ensure_dir,
)

__metaclass__ = type

log = logging.getLogger("assess_autoload_credentials")

def assess_autoload_credentials(juju_bin):
    test_scenarios = [
        ('AWS using environment variables', aws_envvar_test_details),
        ('AWS using credentials file', aws_directory_test_details),
    ]

    for scenario_name, scenario_setup in test_scenarios:
        log.info('* Starting test scenario: {}'.format(scenario_name))
        test_autoload_credentials_stores_details(juju_bin, scenario_setup)

def test_autoload_credentials_stores_details(juju_bin, cloud_details_fn):
    """Test covering loading and storing credentials using autoload-credentials

    :param juju_bin: The full path to the juju binary to use for the test run.
    :param cloud_details_fn: A callable that takes the argument 'user' and
      'tmp_dir' that returns a tuple of:
        (dict -> environment variable changse,
        dict -> expected credential details)
      used to setup creation of credential details & comparison of the result.

    """
    user = 'testing_user'
    with temp_dir() as tmp_dir:
        client = EnvJujuClient.by_version(
            JujuData('local', juju_home=tmp_dir), juju_bin, False)

        env_var_changes, expected_details = cloud_details_fn(user, tmp_dir)
        # Inject well known username.
        env_var_changes.update({'USER': user})

        run_autoload_credentials(client, env_var_changes)

        client.env.load_yaml()

        assert_credentials_contains_expected_results(
            client.env.credentials,
            expected_details)

def assert_credentials_contains_expected_results(credentials, expected):
    if credentials != expected:
        raise ValueError(
            'Actual credentials do not match expected credentials.\n'
            'Expected: {expected}\nGot: {got}\n'.format(
                expected=expected,
                got=credentials
            )
        )

def run_autoload_credentials(client, envvars):
    """Execute the command 'juju autoload-credentials'.

    Simple interaction, calls juju autoload-credentials selects the first
    option and then quits.

    :param client: EnvJujuClient from which juju will be called.
    :param envvars: Dictionary containing environment variables to be used
      during execution.
      Note. Must contain a value for QUESTION_CLOUD_NAME to match against.
      Note. Must contain a value for SAVE_CLOUD_NAME to save against.

    """
    # Get juju path from client as we need to use it interactively.
    process = client.expect(
        'autoload-credentials', extra_env=envvars, include_e=False)
    process.expect(
        '.*1. \w+ credential "{}" \(new\).*'.format(
            envvars['QUESTION_CLOUD_NAME']))
    process.sendline('1')

    process.expect(
        'Enter cloud to which the credential belongs, or Q to quit.*')
    process.sendline(envvars['SAVE_CLOUD_NAME'])
    process.expect(
        'Saved aws credential "{}" to cloud \w+'.format(
            envvars['QUESTION_CLOUD_NAME']))
    process.sendline('q')
    process.expect(pexpect.EOF)

    if process.isalive():
        log.debug('juju process is still running: {}'.format(str(process)))
        process.terminate(force=True)
        raise AssertionError('juju process failed to terminate')

def aws_envvar_test_details(
        user, tmp_dir, access_key='access_key', secret_key='secret_key'):
    env_var_changes = get_aws_environment(user, access_key, secret_key)

    expected_details = get_aws_expected_details_dict(
        user, access_key, secret_key)

    return env_var_changes, expected_details

def aws_directory_test_details(
        user, tmp_dir, access_key='access_key', secret_key='secret_key'
):
    expected_details = get_aws_expected_details_dict(
        'default', access_key, secret_key)

    write_aws_config_file(tmp_dir, access_key, secret_key)

    env_var_changes = dict(
        HOME=tmp_dir,
        SAVE_CLOUD_NAME='aws',
        QUESTION_CLOUD_NAME='default'
    )

    return env_var_changes, expected_details

def get_aws_expected_details_dict(cloud_name, access_key, secret_key):
    # Build credentials yaml file-like datastructure.
    return {
        'credentials': {
            'aws': {
                cloud_name: {
                    'auth-type': 'access-key',
                    'access-key': access_key,
                    'secret-key': secret_key,
                }
            }
        }
    }

def get_aws_environment(user, access_key, secret_key):
    """Return a dictionary containing keys suitable for AWS env vars.

    """
    return dict(
        SAVE_CLOUD_NAME='aws',
        QUESTION_CLOUD_NAME=user,
        AWS_ACCESS_KEY_ID=access_key,
        AWS_SECRET_ACCESS_KEY=secret_key
    )

def write_aws_config_file(tmp_dir, access_key, secret_key):
    """Write aws credentials file to tmp_dir

    :return: String path of created credentials file.

    """
    config_dir = os.path.join(tmp_dir, '.aws')
    config_file = os.path.join(config_dir, 'credentials')
    ensure_dir(config_dir)

    config_contents = dedent("""\
    [default]
    aws_access_key_id={}
    aws_secret_access_key={}
    """.format(access_key, secret_key))

    with open(config_file, 'w') as f:
        f.write(config_contents)

    return config_file

def parse_args(argv):
    parser = argparse.ArgumentParser(description="Test autoload-credentials.")
    parser.add_argument(
        'juju_bin', action=enforce_juju_path,
        help='Full path to the Juju binary.')
    parser.add_argument(
        '--verbose', action='store_const',
        default=logging.INFO, const=logging.DEBUG,
        help='Verbose test harness output.')

    return parser.parse_args(argv)

def main(argv=None):
    args = parse_args(argv)
    configure_logging(args.verbose)

    assess_autoload_credentials(args.juju_bin)
    return 0

if __name__ == '__main__':
    sys.exit(main())
<|endoftext|>


<commit_before># -*- coding: utf-8 -*-
# Natural Language Toolkit: Language Model Unit Tests
#
# Copyright (C) 2001-2019 NLTK Project
# Author: Ilia Kurenkov <ilia.kurenkov@gmail.com>
# URL: <http://nltk.org/>
# For license information, see LICENSE.TXT
"""Smoothing algorithms for language modeling.

According to Chen & Goodman 1995 these should work with both Backoff and
Interpolation.
"""

from nltk.lm.api import Smoothing

def _count_non_zero_vals(dictionary):
    return sum(1.0 for c in dictionary.values() if c > 0)

class WittenBell(Smoothing):
    """Witten-Bell smoothing."""

    def __init__(self, vocabulary, counter, discount=0.1, **kwargs):
        super(WittenBell, self).__init__(vocabulary, counter, *kwargs)
        self.counts = counter

    def alpha_gamma(self, word, context):
        gamma = self.gamma(context)
        return (1.0 - gamma) * self.alpha(word, context), gamma

    def unigram_score(self, word):
        return self.counts.unigrams.freq(word)

    def alpha(self, word, context):
        return self.counts[context].freq(word)

    def gamma(self, context):
        n_plus = _count_non_zero_vals(self.counts[context])
        return n_plus / (n_plus + self.counts[len(context) + 1].N())

class KneserNey(Smoothing):
    """Kneser-Ney Smoothing."""

    def __init__(self, vocabulary, counter, discount=0.1, **kwargs):
        super(KneserNey, self).__init__(vocabulary, counter, *kwargs)
        self.discount = discount
        self.vocabulary = vocabulary

    def unigram_score(self, word):
        return 1.0 / len(self.vocabulary)

    def alpha_gamma(self, word, context):
        prefix_counts = self.counts[context]
        return self.alpha(word, prefix_counts), self.gamma(prefix_counts)

    def alpha(self, word, prefix_counts):
        return max(prefix_counts[word] - self.discount, 0.0) / prefix_counts.N()

    def gamma(self, prefix_counts):
        return self.discount * _count_non_zero_vals(prefix_counts) / prefix_counts.N()

class GoodTuring(Smoothing):
    """Good-Turing Smoothing"""
    def __init__(self, vocabulary, counter, **kwargs):
        super(GoodTuring, self).__init__(vocabulary, counter, *kwargs)
        self.counts = counter
        self.vocabulary = vocabulary

    def unigram_score(self, word):
        word_count = self.counts[word]
        count_plus_1 = 0.
        for everyContext in self.counts.keys():
            if len(everyContext.split()) == word_count+1:
                count_plus_1 += 1
        return count_plus_1 / len(self.vocabulary)
<commit_msg>There's not need to override the super class' self.vocab with self.vocabulary<commit_after># -*- coding: utf-8 -*-
# Natural Language Toolkit: Language Model Unit Tests
#
# Copyright (C) 2001-2019 NLTK Project
# Author: Ilia Kurenkov <ilia.kurenkov@gmail.com>
# URL: <http://nltk.org/>
# For license information, see LICENSE.TXT
"""Smoothing algorithms for language modeling.

According to Chen & Goodman 1995 these should work with both Backoff and
Interpolation.
"""

from nltk.lm.api import Smoothing

def _count_non_zero_vals(dictionary):
    return sum(1.0 for c in dictionary.values() if c > 0)

class WittenBell(Smoothing):
    """Witten-Bell smoothing."""

    def __init__(self, vocabulary, counter, discount=0.1, **kwargs):
        super(WittenBell, self).__init__(vocabulary, counter, *kwargs)
        self.counts = counter

    def alpha_gamma(self, word, context):
        gamma = self.gamma(context)
        return (1.0 - gamma) * self.alpha(word, context), gamma

    def unigram_score(self, word):
        return self.counts.unigrams.freq(word)

    def alpha(self, word, context):
        return self.counts[context].freq(word)

    def gamma(self, context):
        n_plus = _count_non_zero_vals(self.counts[context])
        return n_plus / (n_plus + self.counts[len(context) + 1].N())

class KneserNey(Smoothing):
    """Kneser-Ney Smoothing."""

    def __init__(self, vocabulary, counter, discount=0.1, **kwargs):
        super(KneserNey, self).__init__(vocabulary, counter, *kwargs)
        self.discount = discount

    def unigram_score(self, word):
        return 1.0 / len(self.vocab)

    def alpha_gamma(self, word, context):
        prefix_counts = self.counts[context]
        return self.alpha(word, prefix_counts), self.gamma(prefix_counts)

    def alpha(self, word, prefix_counts):
        return max(prefix_counts[word] - self.discount, 0.0) / prefix_counts.N()

    def gamma(self, prefix_counts):
        return self.discount * _count_non_zero_vals(prefix_counts) / prefix_counts.N()

class GoodTuring(Smoothing):
    """Good-Turing Smoothing"""
    def __init__(self, vocabulary, counter, **kwargs):
        super(GoodTuring, self).__init__(vocabulary, counter, *kwargs)
        self.counts = counter
        self.vocabulary = vocabulary

    def unigram_score(self, word):
        word_count = self.counts[word]
        count_plus_1 = 0.
        for everyContext in self.counts.keys():
            if len(everyContext.split()) == word_count+1:
                count_plus_1 += 1
        return count_plus_1 / len(self.vocabulary)
<|endoftext|>


<commit_before># -*- coding: utf-8 -*-
"""
Disqus static comment plugin for Pelican
====================================
This plugin adds a disqus_comments property to all articles.
Comments are fetched at generation time using disqus API.
"""

from disqusapi import DisqusAPI, Paginator
from pelican import signals

def initialized(pelican):
    from pelican.settings import _DEFAULT_CONFIG
    _DEFAULT_CONFIG.setdefault('DISQUS_SECRET_KEY', '')
    _DEFAULT_CONFIG.setdefault('DISQUS_PUBLIC_KEY', '')
    if pelican:
        pelican.settings.setdefault('DISQUS_SECRET_KEY', '')
        pelican.settings.setdefault('DISQUS_PUBLIC_KEY', '')

def disqus_static(generator):
    disqus = DisqusAPI(generator.settings['DISQUS_SECRET_KEY'],
                       generator.settings['DISQUS_PUBLIC_KEY'])
    # first retrieve the threads
    threads = Paginator(disqus.threads.list,
                        forum=generator.settings['DISQUS_SITENAME'])
    # build a {thread_id: title} dict
    thread_dict = {}
    for thread in threads:
        thread_dict[thread['id']] = thread['title']

    # now retrieve the posts
    posts = Paginator(disqus.posts.list,
                      forum=generator.settings['DISQUS_SITENAME'])

    # build a {post_id: [child_post1, child_post2, ...]} dict
    child_dict = {}
    for post in posts:
        if post['id'] not in child_dict.keys():
            child_dict[post['id']] = []
        if post['parent'] is not None:
            if str(post['parent']) not in child_dict.keys():
                child_dict[str(post['parent'])] = []
            child_dict[str(post['parent'])].append(post)

    # build a {title: [post1, post2, ...]} dict
    post_dict = {}
    for post in posts:
        build_post_dict(post_dict, child_dict, thread_dict, post)

    for article in generator.articles:
        if article.title in post_dict:
            article.disqus_comments = post_dict[article.title]
            article.disqus_comment_count = sum([
                postcounter(post) for post in post_dict[article.title]])

def postcounter(node):
    return 1 + sum([postcounter(n) for n in node['children']])

def build_post_dict(post_dict, child_dict, thread_dict, post):
    if post['thread'] not in thread_dict.keys():
        return # invalid thread, should never happen

    build_child_dict(child_dict, post)

    if post['parent'] is not None:
        return # this is a child post, don't want to display it here

    if thread_dict[post['thread']] not in post_dict.keys():
        post_dict[thread_dict[post['thread']]] = []
    post_dict[thread_dict[post['thread']]].append(post)

def build_child_dict(child_dict, post):
    post['children'] = child_dict[post['id']]
    for child in child_dict[post['id']]:
        build_child_dict(child_dict, child)

def register():
    signals.initialized.connect(initialized)
    signals.article_generator_finalized.connect(disqus_static)
<commit_msg>Update disqus_static.py<commit_after># -*- coding: utf-8 -*-
"""
Disqus static comment plugin for Pelican
====================================
This plugin adds a disqus_comments property to all articles.
Comments are fetched at generation time using disqus API.
"""

from disqusapi import DisqusAPI, Paginator
from pelican import signals

def initialized(pelican):
    from pelican.settings import DEFAULT_CONFIG
    DEFAULT_CONFIG.setdefault('DISQUS_SECRET_KEY', '')
    DEFAULT_CONFIG.setdefault('DISQUS_PUBLIC_KEY', '')
    if pelican:
        pelican.settings.setdefault('DISQUS_SECRET_KEY', '')
        pelican.settings.setdefault('DISQUS_PUBLIC_KEY', '')

def disqus_static(generator):
    disqus = DisqusAPI(generator.settings['DISQUS_SECRET_KEY'],
                       generator.settings['DISQUS_PUBLIC_KEY'])
    # first retrieve the threads
    threads = Paginator(disqus.threads.list,
                        forum=generator.settings['DISQUS_SITENAME'])
    # build a {thread_id: title} dict
    thread_dict = {}
    for thread in threads:
        thread_dict[thread['id']] = thread['title']

    # now retrieve the posts
    posts = Paginator(disqus.posts.list,
                      forum=generator.settings['DISQUS_SITENAME'])

    # build a {post_id: [child_post1, child_post2, ...]} dict
    child_dict = {}
    for post in posts:
        if post['id'] not in child_dict.keys():
            child_dict[post['id']] = []
        if post['parent'] is not None:
            if str(post['parent']) not in child_dict.keys():
                child_dict[str(post['parent'])] = []
            child_dict[str(post['parent'])].append(post)

    # build a {title: [post1, post2, ...]} dict
    post_dict = {}
    for post in posts:
        build_post_dict(post_dict, child_dict, thread_dict, post)

    for article in generator.articles:
        if article.title in post_dict:
            article.disqus_comments = post_dict[article.title]
            article.disqus_comment_count = sum([
                postcounter(post) for post in post_dict[article.title]])

def postcounter(node):
    return 1 + sum([postcounter(n) for n in node['children']])

def build_post_dict(post_dict, child_dict, thread_dict, post):
    if post['thread'] not in thread_dict.keys():
        return # invalid thread, should never happen

    build_child_dict(child_dict, post)

    if post['parent'] is not None:
        return # this is a child post, don't want to display it here

    if thread_dict[post['thread']] not in post_dict.keys():
        post_dict[thread_dict[post['thread']]] = []
    post_dict[thread_dict[post['thread']]].append(post)

def build_child_dict(child_dict, post):
    post['children'] = child_dict[post['id']]
    for child in child_dict[post['id']]:
        build_child_dict(child_dict, child)

def register():
    signals.initialized.connect(initialized)
    signals.article_generator_finalized.connect(disqus_static)
<|endoftext|>