<commit_before>import threading
from collections import OrderedDict

from django import forms
from django.conf import settings
from django.core.cache import cache
from django.core.exceptions import ValidationError
from django.db import models
from django.utils.translation import ugettext_lazy as _

from c3nav.mapdata.fields import JSONField
from c3nav.mapdata.models import MapUpdate, WayType

class RouteOptions(models.Model):
    user = models.OneToOneField(settings.AUTH_USER_MODEL, on_delete=models.CASCADE, primary_key=True)
    data = JSONField(default={})

    class Meta:
        verbose_name = _('Route options')
        verbose_name_plural = _('Route options')
        default_related_name = 'routeoptions'

    fields_cached = None
    fields_cache_key = None
    fields_cache_lock = threading.Lock()

    @classmethod
    def build_fields(cls):
        fields = OrderedDict()
        fields['mode'] = forms.ChoiceField(
            label=_('Routing mode'),
            choices=(('fastest', _('fastest')), ('shortest', _('shortest'))),
            initial='fastest'
        )
        fields['walk_speed'] = forms.ChoiceField(
            label=_('Walk speed'),
            choices=(('slow', _('slow')), ('default', _('default')), ('fast', _('fast'))),
            initial='default'
        )

        for waytype in WayType.objects.all():
            choices = []
            choices.append(('allow', _('allow')))
            if waytype.up_separate:
                choices.append(('avoid_up', _('avoid upwards')))
                choices.append(('avoid_down', _('avoid downwards')))
                choices.append(('avoid', _('avoid completely')))
            else:
                choices.append(('avoid', _('avoid')))
            fields['waytype_%d' % waytype.pk] = forms.ChoiceField(
                label=waytype.title_plural,
                choices=tuple(choices),
                initial='allow'
            )

        return fields

    @classmethod
    def get_fields(cls):
        cache_key = MapUpdate.current_cache_key()
        if cls.fields_cache_key != cache_key:
            with cls.fields_cache_lock:
                cls.fields_cache_key = cache_key
                cls.fields_cached = cls.build_fields()
        return cls.fields_cached

    @staticmethod
    def get_cache_key(pk):
        return 'routing:options:user:%d' % pk

    @classmethod
    def get_for_user(cls, user):
        cache_key = cls.get_cache_key(user.pk)
        result = cache.get(cache_key, None)
        if result:
            return result

        try:
            result = user.routeoptions
        except AttributeError:
            result = None

        if result:
            cache.set(cache_key, result, 900)

        return result

    @classmethod
    def get_for_request(cls, request):
        session_options = request.session.get('route_options', None)
        user_options = None
        if request.user.is_authenticated:
            user_options = cls.get_for_user(request.user)

        if session_options and not user_options:
            user_options = session_options
            user_options.user = request.user
            user_options.save()
            request.session.pop('session_options')

        options = cls(request=request)
        options.update(user_options or session_options or {}, ignore_errors=True)
        return options

    def __init__(self, *args, request=None, **kwargs):
        super().__init__(*args, **kwargs)
        new_data = OrderedDict()
        for name, field in self.get_fields().items():
            value = self.data.get(name)
            if value is None or value not in dict(field.choices):
                value = field.initial
            new_data[name] = value
        self.data = new_data

        self.request = request

    def __getitem__(self, key):
        try:
            return self.data[key]
        except AttributeError:
            return self.get_fields()[key].initial

    def update(self, value_dict, ignore_errors=False, ignore_unknown=False):
        if not value_dict:
            return
        if isinstance(value_dict, RouteOptions):
            value_dict = value_dict.data
        fields = self.get_fields()
        for key, value in value_dict.items():
            field = fields.get(key)
            if not field:
                if ignore_errors or ignore_unknown:
                    continue
                raise ValidationError(_('Unknown route option: %s') % key)
            if value is None or value not in dict(field.choices):
                if ignore_errors:
                    continue
                raise ValidationError(_('Invalid value for route option %s.') % key)
            self.data[key] = value

    def __setitem__(self, key, value):
        self.update({key: value})

    @property
    def walk_factor(self):
        return {'slow': 0.8, 'default': 1, 'fast': 1.2}[self['walk_speed']]

    def get(self, key, default):
        try:
            return self[key]
        except AttributeError:
            return default

    def serialize(self):
        return [
            {
                'name': name,
                'type': field.widget.input_type,
                'label': field.label,
                'choices': [
                    {
                        'name': choice_name,
                        'title': choice_title,
                    }
                    for choice_name, choice_title in field.choices
                ],
                'value': self[name],
            }
            for name, field in self.get_fields().items()
        ]

    def save(self, *args, **kwargs):
        if self.request is None or self.request.user.is_authenticated:
            self.user = self.request.user
            return super().save(*args, **kwargs)

        self.request.session['route_options'] = self

    def items(self):
        yield from self.data.items()
<commit_msg>fix saving route options<commit_after>import threading
from collections import OrderedDict

from django import forms
from django.conf import settings
from django.core.cache import cache
from django.core.exceptions import ValidationError
from django.db import models
from django.utils.translation import ugettext_lazy as _

from c3nav.mapdata.fields import JSONField
from c3nav.mapdata.models import MapUpdate, WayType

class RouteOptions(models.Model):
    user = models.OneToOneField(settings.AUTH_USER_MODEL, on_delete=models.CASCADE, primary_key=True)
    data = JSONField(default={})

    class Meta:
        verbose_name = _('Route options')
        verbose_name_plural = _('Route options')
        default_related_name = 'routeoptions'

    fields_cached = None
    fields_cache_key = None
    fields_cache_lock = threading.Lock()

    @classmethod
    def build_fields(cls):
        fields = OrderedDict()
        fields['mode'] = forms.ChoiceField(
            label=_('Routing mode'),
            choices=(('fastest', _('fastest')), ('shortest', _('shortest'))),
            initial='fastest'
        )
        fields['walk_speed'] = forms.ChoiceField(
            label=_('Walk speed'),
            choices=(('slow', _('slow')), ('default', _('default')), ('fast', _('fast'))),
            initial='default'
        )

        for waytype in WayType.objects.all():
            choices = []
            choices.append(('allow', _('allow')))
            if waytype.up_separate:
                choices.append(('avoid_up', _('avoid upwards')))
                choices.append(('avoid_down', _('avoid downwards')))
                choices.append(('avoid', _('avoid completely')))
            else:
                choices.append(('avoid', _('avoid')))
            fields['waytype_%d' % waytype.pk] = forms.ChoiceField(
                label=waytype.title_plural,
                choices=tuple(choices),
                initial='allow'
            )

        return fields

    @classmethod
    def get_fields(cls):
        cache_key = MapUpdate.current_cache_key()
        if cls.fields_cache_key != cache_key:
            with cls.fields_cache_lock:
                cls.fields_cache_key = cache_key
                cls.fields_cached = cls.build_fields()
        return cls.fields_cached

    @staticmethod
    def get_cache_key(pk):
        return 'routing:options:user:%d' % pk

    @classmethod
    def get_for_user(cls, user):
        cache_key = cls.get_cache_key(user.pk)
        result = cache.get(cache_key, None)
        if result:
            return result

        try:
            result = user.routeoptions
        except AttributeError:
            result = None

        if result:
            cache.set(cache_key, result, 900)

        return result

    @classmethod
    def get_for_request(cls, request):
        session_options = request.session.get('route_options', None)
        if session_options is not None:
            session_options = cls(request=request)
            session_options.update(session_options, ignore_errors=True)

        user_options = None
        if request.user.is_authenticated:
            user_options = cls.get_for_user(request.user)
            user_options.clean_data()

        if session_options and not user_options:
            user_options = session_options
            user_options.user = request.user
            user_options.save()
            request.session.pop('session_options')

        return user_options or session_options or cls(request=request)

    def clean_data(self):
        new_data = OrderedDict()
        for name, field in self.get_fields().items():
            value = self.data.get(name)
            if value is None or value not in dict(field.choices):
                value = field.initial
            new_data[name] = value
        self.data = new_data

    def __init__(self, *args, request=None, **kwargs):
        super().__init__(*args, **kwargs)
        self.clean_data()

        self.request = request

    def __getitem__(self, key):
        try:
            return self.data[key]
        except AttributeError:
            return self.get_fields()[key].initial

    def update(self, value_dict, ignore_errors=False, ignore_unknown=False):
        if not value_dict:
            return
        if isinstance(value_dict, RouteOptions):
            value_dict = value_dict.data
        fields = self.get_fields()
        for key, value in value_dict.items():
            field = fields.get(key)
            if not field:
                if ignore_errors or ignore_unknown:
                    continue
                raise ValidationError(_('Unknown route option: %s') % key)
            if value is None or value not in dict(field.choices):
                if ignore_errors:
                    continue
                raise ValidationError(_('Invalid value for route option %s.') % key)
            self.data[key] = value

    def __setitem__(self, key, value):
        self.update({key: value})

    @property
    def walk_factor(self):
        return {'slow': 0.8, 'default': 1, 'fast': 1.2}[self['walk_speed']]

    def get(self, key, default):
        try:
            return self[key]
        except AttributeError:
            return default

    def serialize(self):
        return [
            {
                'name': name,
                'type': field.widget.input_type,
                'label': field.label,
                'choices': [
                    {
                        'name': choice_name,
                        'title': choice_title,
                    }
                    for choice_name, choice_title in field.choices
                ],
                'value': self[name],
            }
            for name, field in self.get_fields().items()
        ]

    def save(self, *args, **kwargs):
        if self.request is None or self.request.user.is_authenticated:
            self.user = self.request.user
            return super().save(*args, **kwargs)

        self.request.session['route_options'] = self.data

    def items(self):
        yield from self.data.items()
<|endoftext|>


<commit_before># (c) 2012, Michael DeHaan <michael.dehaan@gmail.com>
#
# This file is part of Ansible
#
# Ansible is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# Ansible 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 General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with Ansible.  If not, see <http://www.gnu.org/licenses/>.

class Group(object):
    ''' a group of ansible hosts '''

    __slots__ = [ 'name', 'hosts', 'vars', 'child_groups', 'parent_groups', 'depth' ]

    def __init__(self, name=None):

        self.depth = 0
        self.name = name
        self.hosts = []
        self.vars = {}
        self.child_groups = []
        self.parent_groups = []
        if self.name is None:
            raise Exception("group name is required")

    def add_child_group(self, group):

        if self == group:
            raise Exception("can't add group to itself")
        self.child_groups.append(group)
        group.depth = group.depth + 1
        group.parent_groups.append(self)

    def add_host(self, host):

        self.hosts.append(host)
        host.add_group(self)

    def set_variable(self, key, value):

        self.vars[key] = value

    def get_hosts(self):

        hosts = set()
        for kid in self.child_groups:
            hosts.update(kid.get_hosts())
        hosts.update(self.hosts)
        return list(hosts)

    def get_variables(self):

        vars = {}
        # FIXME: verify this variable override order is what we want
        for ancestor in self.get_ancestors():
            vars.update(ancestor.get_variables())
        vars.update(self.vars)
        return vars

    def _get_ancestors(self):

        results = {}
        for g in self.parent_groups:
            results[g.name] = g
            results.update(g._get_ancestors())
        return results

    def get_ancestors(self):

        return self._get_ancestors().values()

<commit_msg>Group should only return its own variables<commit_after># (c) 2012, Michael DeHaan <michael.dehaan@gmail.com>
#
# This file is part of Ansible
#
# Ansible is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# Ansible 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 General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with Ansible.  If not, see <http://www.gnu.org/licenses/>.

class Group(object):
    ''' a group of ansible hosts '''

    __slots__ = [ 'name', 'hosts', 'vars', 'child_groups', 'parent_groups', 'depth' ]

    def __init__(self, name=None):

        self.depth = 0
        self.name = name
        self.hosts = []
        self.vars = {}
        self.child_groups = []
        self.parent_groups = []
        if self.name is None:
            raise Exception("group name is required")

    def add_child_group(self, group):

        if self == group:
            raise Exception("can't add group to itself")
        self.child_groups.append(group)
        group.depth = group.depth + 1
        group.parent_groups.append(self)

    def add_host(self, host):

        self.hosts.append(host)
        host.add_group(self)

    def set_variable(self, key, value):

        self.vars[key] = value

    def get_hosts(self):

        hosts = set()
        for kid in self.child_groups:
            hosts.update(kid.get_hosts())
        hosts.update(self.hosts)
        return list(hosts)

    def get_variables(self):
        return self.vars.copy()

    def _get_ancestors(self):

        results = {}
        for g in self.parent_groups:
            results[g.name] = g
            results.update(g._get_ancestors())
        return results

    def get_ancestors(self):

        return self._get_ancestors().values()

<|endoftext|>


<commit_before>#!/usr/bin/env Python
"""
Copyright (c) 2013 The Regents of the University of California, AMERICAN INSTITUTES FOR RESEARCH
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.
"""
"""
@author Gabe Fierro gt.fierro@berkeley.edu github.com/gtfierro
"""
"""
Performs a basic assignee disambiguation
"""
from collections import defaultdict, deque
import uuid
from string import lowercase as alphabet
import re
import md5
import cPickle as pickle
import alchemy
from collections import Counter
from Levenshtein import jaro_winkler
from alchemy import get_config, match
from alchemy.schema import *
from alchemy.match import commit_inserts, commit_updates
from handlers.xml_util import normalize_utf8
from datetime import datetime
from sqlalchemy.sql import or_
from sqlalchemy.sql.expression import bindparam
from unidecode import unidecode
from tasks import bulk_commit_inserts, bulk_commit_updates
import multiprocessing
import itertools
import sys

config = get_config()

THRESHOLD = config.get("assignee").get("threshold")

uuid_to_object = {}
uuid_to_cleanid = {}
letter_to_cleanid = {}
uuids_by_cleanidletter = defaultdict(list)

grant_uuids = set()
app_uuids = set()
grantsessiongen = alchemy.session_generator(dbtype='grant')
appsessiongen = alchemy.session_generator(dbtype='application')

nodigits = re.compile(r'[a-z ]')
stoplist = ['the','of','and','a','an','at']

def isgrant(obj):
    """
    returns True of obj is from Grant table, False if from App table
    """
    return hasattr(obj, 'patent')

def get_cleanid(obj):
    """
    Returns a cleaned string version of the object representation:

    if obj has an organization, uses that. If obj has first/last name,
    uses "firstname|lastname"

    Changes severything to lowercase and removes everything that isn't [a-z ]
    """
    cleanid = ''
    if obj.organization:
        cleanid = obj.organization
    else:
        try:
            cleanid = obj.name_first + obj.name_last
        except:
            cleanid = ''
    cleanid = cleanid.lower()
    cleanid = ' '.join(filter(lambda x:
                        x not in stoplist,
                        cleanid.split()))
    cleanid = ''.join(nodigits.findall(cleanid)).strip()
    return cleanid

def get_similarity(uuid1, uuid2):
    clean1 = uuid_to_cleanid[uuid1]
    clean2 = uuid_to_cleanid[uuid2]
    if clean1 == clean2:
      return 1.0
    return jaro_winkler(clean1,clean2,0.0)

def disambiguate_letter(letter):
    groups = defaultdict(list)
    bucket = uuids_by_cleanidletter[letter]
    uuidsremaining = bucket[:]
    groupkeys = []
    print len(bucket),'raw assignees for letter:', letter
    i = 1
    while True:
        if not uuidsremaining:
            break
        i += 1
        uuid = uuidsremaining.pop()
        if i%10000 == 0:
            print i, datetime.now()
        matcheduuid = False
        for groupkey in groupkeys:
            if get_similarity(uuid, groupkey) >= THRESHOLD:
                groups[groupkey].append(uuid)
                matcheduuid = True
                break
        if matcheduuid: continue
        groups[uuid].append(uuid)
        groupkeys.append(uuid)
    return groups

def create_disambiguated_record_for_block(block):
    grant_assignee_inserts = []
    app_assignee_inserts = []
    patentassignee_inserts = []
    applicationassignee_inserts = []
    grant_rawassignee_updates = []
    app_rawassignee_updates = []
    ra_objs = [uuid_to_object[uuid] for uuid in block]
    # vote on the disambiguated assignee parameters
    freq = defaultdict(Counter)
    param = {}

    for ra in ra_objs:
        for k,v in ra.summarize.items():
            if not v:
                v = ''
            param[k] = v

    if not param.has_key('organization'):
        param['organization'] = ''
    if not param.has_key('type'):
        param['type'] = ''
    if not param.has_key('name_last'):
        param['name_last'] = ''
    if not param.has_key('name_first'):
        param['name_first'] = ''
    if not param.has_key('residence'):
        param['residence'] = ''
    if not param.has_key('nationality'):
        param['nationality'] = ''
    if param.has_key('type'):
        if not param['type'].isdigit():
            param['type'] = ''

    # create persistent identifier
    if param["organization"]:
        param["id"] = md5.md5(unidecode(param["organization"])).hexdigest()
    elif param["name_last"]:
        param["id"] = md5.md5(unidecode(param["name_last"]+param["name_first"])).hexdigest()
    else:
        param["id"] = md5.md5('').hexdigest()

    grant_assignee_inserts.append(param)
    app_assignee_inserts.append(param)

    # inserts for patent_assignee and appliation_assignee tables
    patents = filter(lambda x: x, map(lambda x: getattr(x,'patent_id',None), ra_objs))
    patentassignee_inserts.extend({'patent_id': x, 'assignee_id': param['id']} for x in patents)
    applications = filter(lambda x: x, map(lambda x: getattr(x,'application_id',None), ra_objs))
    applicationassignee_inserts.extend([{'application_id': x, 'assignee_id': param['id']} for x in applications])

    # update statements for rawassignee tables
    for ra in ra_objs:
        if isgrant(ra):
            grant_rawassignee_updates.append({'pk': ra.uuid, 'update': param['id']})
        else:
            app_rawassignee_updates.append({'pk': ra.uuid, 'update': param['id']})
    return grant_assignee_inserts, app_assignee_inserts, patentassignee_inserts, applicationassignee_inserts, grant_rawassignee_updates, app_rawassignee_updates

def run_disambiguation():
    """
    Runs disambiguation algorithm on grant and application assignees from
    the database indicated by lib/alchemy/config
    """
    # retrieve database connections and pull in all assignees from
    # both grant and application databases
    grtsesh = grantsessiongen()
    appsesh = appsessiongen()
    print 'fetching raw assignees',datetime.now()
    rawassignees = list(grtsesh.query(RawAssignee))
    rawassignees.extend(list(appsesh.query(App_RawAssignee)))
    # clear the destination tables
    if alchemy.is_mysql():
        grtsesh.execute('truncate assignee; truncate patent_assignee;')
        appsesh.execute('truncate assignee; truncate application_assignee;')
    else:
        grtsesh.execute('delete from assignee; delete from patent_assignee;')
        appsesh.execute('delete from assignee; delete from patent_assignee;')
    print 'cleaning ids', datetime.now()
    # uses the get_cleanid method to remove undesirable characters and
    # normalize to case and group by first letter
    for ra in rawassignees:
        uuid_to_object[ra.uuid] = ra
        cleanid = get_cleanid(ra)
        uuid_to_cleanid[ra.uuid] = cleanid
        if not cleanid:
          continue
        firstletter = cleanid[0]
        uuids_by_cleanidletter[firstletter].append(ra.uuid)

    print 'disambiguating blocks', datetime.now()
    # disambiguates each of the letter blocks using
    # the list of assignees as a stack and only performing
    # jaro-winkler comparisons on the first item of each block
    allrecords = []
    for letter in alphabet:
        print 'disambiguating','({0})'.format(letter),datetime.now()
        lettergroup = disambiguate_letter(letter)
        print 'got',len(lettergroup),'records'
        print 'creating disambiguated records','({0})'.format(letter),datetime.now()
        allrecords.extend(lettergroup.values())
    # create the attributes for the disambiguated assignee record from the
    # raw records placed into a block in the disambiguation phase
    res = map(create_disambiguated_record_for_block, allrecords)
    mid = itertools.izip(*res)
    grant_assignee_inserts = list(itertools.chain.from_iterable(mid.next()))
    app_assignee_inserts = list(itertools.chain.from_iterable(mid.next()))
    patentassignee_inserts = list(itertools.chain.from_iterable(mid.next()))
    applicationassignee_inserts = list(itertools.chain.from_iterable(mid.next()))
    grant_rawassignee_updates = list(itertools.chain.from_iterable(mid.next()))
    app_rawassignee_updates = list(itertools.chain.from_iterable(mid.next()))

    # write out the insert counts for each table into a text file
    with open('mid.txt','wb') as f:
      f.write(str(len(grant_assignee_inserts))+'\n')
      f.write(str(len(app_assignee_inserts))+'\n')
      f.write(str(len(patentassignee_inserts))+'\n')
      f.write(str(len(applicationassignee_inserts))+'\n')
      f.write(str(len(grant_rawassignee_updates))+'\n')
      f.write(str(len(app_rawassignee_updates))+'\n')
    # insert disambiguated assignee records
    bulk_commit_inserts(grant_assignee_inserts, Assignee.__table__, alchemy.is_mysql(), 20000, 'grant')
    bulk_commit_inserts(app_assignee_inserts, App_Assignee.__table__, alchemy.is_mysql(), 20000, 'application')
    # insert patent/assignee link records
    bulk_commit_inserts(patentassignee_inserts, patentassignee, alchemy.is_mysql(), 20000, 'grant')
    bulk_commit_inserts(applicationassignee_inserts, applicationassignee, alchemy.is_mysql(), 20000, 'application')
    # update rawassignees with their disambiguated record
    bulk_commit_updates('assignee_id', grant_rawassignee_updates, RawAssignee.__table__, alchemy.is_mysql(), 20000, 'grant')
    bulk_commit_updates('assignee_id', app_rawassignee_updates, App_RawAssignee.__table__, alchemy.is_mysql(), 20000, 'application')

if __name__=='__main__':
  run_disambiguation()
<commit_msg>add nber substitutions to assignee disambig<commit_after>#!/usr/bin/env Python
"""
Copyright (c) 2013 The Regents of the University of California, AMERICAN INSTITUTES FOR RESEARCH
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.
"""
"""
@author Gabe Fierro gt.fierro@berkeley.edu github.com/gtfierro
"""
"""
Performs a basic assignee disambiguation
"""
from collections import defaultdict, deque
import uuid
from string import lowercase as alphabet
import re
import md5
import cPickle as pickle
import alchemy
from collections import Counter
from Levenshtein import jaro_winkler
from alchemy import get_config, match
from alchemy.schema import *
from alchemy.match import commit_inserts, commit_updates
from handlers.xml_util import normalize_utf8
from datetime import datetime
from sqlalchemy.sql import or_
from sqlalchemy.sql.expression import bindparam
from unidecode import unidecode
from tasks import bulk_commit_inserts, bulk_commit_updates
import multiprocessing
import itertools
import sys
import json

config = get_config()

THRESHOLD = config.get("assignee").get("threshold")

uuid_to_object = {}
uuid_to_cleanid = {}
letter_to_cleanid = {}
uuids_by_cleanidletter = defaultdict(list)

grant_uuids = set()
app_uuids = set()
grantsessiongen = alchemy.session_generator(dbtype='grant')
appsessiongen = alchemy.session_generator(dbtype='application')

nodigits = re.compile(r'[a-z ]')
stoplist = ['the','of','and','a','an','at']
substitutions = json.load(open('nber_substitutions.json'))

def isgrant(obj):
    """
    returns True of obj is from Grant table, False if from App table
    """
    return hasattr(obj, 'patent')

def get_cleanid(obj):
    """
    Returns a cleaned string version of the object representation:

    if obj has an organization, uses that. If obj has first/last name,
    uses "firstname|lastname"

    Changes severything to lowercase and removes everything that isn't [a-z ]
    """
    cleanid = ''
    if obj.organization:
        cleanid = obj.organization
    else:
        try:
            cleanid = obj.name_first + obj.name_last
        except:
            cleanid = ''
    cleanid = cleanid.lower()
    cleanid = ' '.join(filter(lambda x:
                        x not in stoplist,
                        cleanid.split()))
    cleanid = ''.join(nodigits.findall(cleanid)).strip()
    for pair in substitutions:
      cleanid = cleanid.replace(pair[0], pair[1])
    return cleanid

def get_similarity(uuid1, uuid2):
    clean1 = uuid_to_cleanid[uuid1]
    clean2 = uuid_to_cleanid[uuid2]
    if clean1 == clean2:
      return 1.0
    return jaro_winkler(clean1,clean2,0.0)

def disambiguate_letter(letter):
    groups = defaultdict(list)
    bucket = uuids_by_cleanidletter[letter]
    uuidsremaining = bucket[:]
    groupkeys = []
    print len(bucket),'raw assignees for letter:', letter
    i = 1
    while True:
        if not uuidsremaining:
            break
        i += 1
        uuid = uuidsremaining.pop()
        if i%10000 == 0:
            print i, datetime.now()
        matcheduuid = False
        for groupkey in groupkeys:
            if get_similarity(uuid, groupkey) >= THRESHOLD:
                groups[groupkey].append(uuid)
                matcheduuid = True
                break
        if matcheduuid: continue
        groups[uuid].append(uuid)
        groupkeys.append(uuid)
    return groups

def create_disambiguated_record_for_block(block):
    grant_assignee_inserts = []
    app_assignee_inserts = []
    patentassignee_inserts = []
    applicationassignee_inserts = []
    grant_rawassignee_updates = []
    app_rawassignee_updates = []
    ra_objs = [uuid_to_object[uuid] for uuid in block]
    # vote on the disambiguated assignee parameters
    freq = defaultdict(Counter)
    param = {}

    for ra in ra_objs:
        for k,v in ra.summarize.items():
            if not v:
                v = ''
            param[k] = v

    if not param.has_key('organization'):
        param['organization'] = ''
    if not param.has_key('type'):
        param['type'] = ''
    if not param.has_key('name_last'):
        param['name_last'] = ''
    if not param.has_key('name_first'):
        param['name_first'] = ''
    if not param.has_key('residence'):
        param['residence'] = ''
    if not param.has_key('nationality'):
        param['nationality'] = ''
    if param.has_key('type'):
        if not param['type'].isdigit():
            param['type'] = ''

    # create persistent identifier
    if param["organization"]:
        param["id"] = md5.md5(unidecode(param["organization"])).hexdigest()
    elif param["name_last"]:
        param["id"] = md5.md5(unidecode(param["name_last"]+param["name_first"])).hexdigest()
    else:
        param["id"] = md5.md5('').hexdigest()

    grant_assignee_inserts.append(param)
    app_assignee_inserts.append(param)

    # inserts for patent_assignee and appliation_assignee tables
    patents = filter(lambda x: x, map(lambda x: getattr(x,'patent_id',None), ra_objs))
    patentassignee_inserts.extend({'patent_id': x, 'assignee_id': param['id']} for x in patents)
    applications = filter(lambda x: x, map(lambda x: getattr(x,'application_id',None), ra_objs))
    applicationassignee_inserts.extend([{'application_id': x, 'assignee_id': param['id']} for x in applications])

    # update statements for rawassignee tables
    for ra in ra_objs:
        if isgrant(ra):
            grant_rawassignee_updates.append({'pk': ra.uuid, 'update': param['id']})
        else:
            app_rawassignee_updates.append({'pk': ra.uuid, 'update': param['id']})
    return grant_assignee_inserts, app_assignee_inserts, patentassignee_inserts, applicationassignee_inserts, grant_rawassignee_updates, app_rawassignee_updates

def run_disambiguation():
    """
    Runs disambiguation algorithm on grant and application assignees from
    the database indicated by lib/alchemy/config
    """
    # retrieve database connections and pull in all assignees from
    # both grant and application databases
    grtsesh = grantsessiongen()
    appsesh = appsessiongen()
    print 'fetching raw assignees',datetime.now()
    rawassignees = list(grtsesh.query(RawAssignee))
    rawassignees.extend(list(appsesh.query(App_RawAssignee)))
    # clear the destination tables
    if alchemy.is_mysql():
        grtsesh.execute('truncate assignee; truncate patent_assignee;')
        appsesh.execute('truncate assignee; truncate application_assignee;')
    else:
        grtsesh.execute('delete from assignee; delete from patent_assignee;')
        appsesh.execute('delete from assignee; delete from patent_assignee;')
    print 'cleaning ids', datetime.now()
    # uses the get_cleanid method to remove undesirable characters and
    # normalize to case and group by first letter
    for ra in rawassignees:
        uuid_to_object[ra.uuid] = ra
        cleanid = get_cleanid(ra)
        uuid_to_cleanid[ra.uuid] = cleanid
        if not cleanid:
          continue
        firstletter = cleanid[0]
        uuids_by_cleanidletter[firstletter].append(ra.uuid)

    print 'disambiguating blocks', datetime.now()
    # disambiguates each of the letter blocks using
    # the list of assignees as a stack and only performing
    # jaro-winkler comparisons on the first item of each block
    allrecords = []
    for letter in alphabet:
        print 'disambiguating','({0})'.format(letter),datetime.now()
        lettergroup = disambiguate_letter(letter)
        print 'got',len(lettergroup),'records'
        print 'creating disambiguated records','({0})'.format(letter),datetime.now()
        allrecords.extend(lettergroup.values())
    # create the attributes for the disambiguated assignee record from the
    # raw records placed into a block in the disambiguation phase
    res = map(create_disambiguated_record_for_block, allrecords)
    mid = itertools.izip(*res)
    grant_assignee_inserts = list(itertools.chain.from_iterable(mid.next()))
    app_assignee_inserts = list(itertools.chain.from_iterable(mid.next()))
    patentassignee_inserts = list(itertools.chain.from_iterable(mid.next()))
    applicationassignee_inserts = list(itertools.chain.from_iterable(mid.next()))
    grant_rawassignee_updates = list(itertools.chain.from_iterable(mid.next()))
    app_rawassignee_updates = list(itertools.chain.from_iterable(mid.next()))

    # write out the insert counts for each table into a text file
    with open('mid.txt','wb') as f:
      f.write(str(len(grant_assignee_inserts))+'\n')
      f.write(str(len(app_assignee_inserts))+'\n')
      f.write(str(len(patentassignee_inserts))+'\n')
      f.write(str(len(applicationassignee_inserts))+'\n')
      f.write(str(len(grant_rawassignee_updates))+'\n')
      f.write(str(len(app_rawassignee_updates))+'\n')
    # insert disambiguated assignee records
    bulk_commit_inserts(grant_assignee_inserts, Assignee.__table__, alchemy.is_mysql(), 20000, 'grant')
    bulk_commit_inserts(app_assignee_inserts, App_Assignee.__table__, alchemy.is_mysql(), 20000, 'application')
    # insert patent/assignee link records
    bulk_commit_inserts(patentassignee_inserts, patentassignee, alchemy.is_mysql(), 20000, 'grant')
    bulk_commit_inserts(applicationassignee_inserts, applicationassignee, alchemy.is_mysql(), 20000, 'application')
    # update rawassignees with their disambiguated record
    bulk_commit_updates('assignee_id', grant_rawassignee_updates, RawAssignee.__table__, alchemy.is_mysql(), 20000, 'grant')
    bulk_commit_updates('assignee_id', app_rawassignee_updates, App_RawAssignee.__table__, alchemy.is_mysql(), 20000, 'application')

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


<commit_before>#!/usr/bin/python
import argparse
import sys
import os
import os.path
import io
import time
import traceback
import ctypes
import subprocess
from subprocess import Popen, PIPE
import shutil
from optparse import OptionParser
import requests
from requests_toolbelt import MultipartEncoder
import urllib
import urllib2
import json
import tarfile
import zipfile
import glob

try:
    from cStringIO import StringIO
except:
    from StringIO import StringIO

from biokbase.workspace.client import Workspace

BUF_SIZE = 8*1024 # default HTTP LIB client buffer_size

# Base class for Transform service
class TransformBase:

    def __init__(self, args):
        self.shock_url = args.shock_url
        self.inobj_id = args.inobj_id
        self.sdir = args.sdir
        self.itmp = args.itmp
        self.otmp = args.otmp
        self.token = os.environ.get('KB_AUTH_TOKEN')
        self.ssl_verify = True

    def upload_to_shock(self):
        if self.token is None:
            raise Exception("Unable to find token!")
        filePath = "{}/{}".format(self.sdir,self.otmp)

        #build the header
        header = dict()
        header["Authorization"] = "Oauth %s" % self.token

        dataFile = open(os.path.abspath(filePath))
        m = MultipartEncoder(fields={'upload': (os.path.split(filePath)[-1], dataFile)})
        header['Content-Type'] = m.content_type

        try:
            response = requests.post(self.shock_url + "/node", headers=header, data=m, allow_redirects=True, verify=self.ssl_verify)
            dataFile.close()

            if not response.ok:
                response.raise_for_status()

            result = response.json()

            if result['error']:
                raise Exception(result['error'][0])
            else:
                return result["data"]
        except:
            dataFile.close()
            raise

    def download_shock_data(self) :
        if self.token is None:
            raise Exception("Unable to find token!")
        # TODO: Improve folder checking
        try:
            os.mkdir(self.sdir)
        except:
            pass

        inids = self.inobj_id.split(',')

        for i in range(len(inids)):
            meta_req = urllib2.Request("{}/node/{}".format(self.shock_url, inids[i]))
            meta_req.add_header('Authorization',"OAuth {}".format(self.token))
            data_req = urllib2.Request("{}/node/{}?download_raw".format(self.shock_url, inids[i]))
            data_req.add_header('Authorization',"OAuth {}".format(self.token))

            meta = urllib2.urlopen(meta_req)
            md = json.loads(meta.read())
            meta.close()

            rdata = urllib2.urlopen(data_req)

            data = StringIO(rdata.read())
            magic = data.read(4)
            data.seek(0)

            if magic.startswith('\x1f\x8b') or magic.startswith('\x42\x5a') : # gz or bz
                my_tar = tarfile.open(fileobj=data, mode="r|*", bufsize=BUF_SIZE)
                if len(inids) > 1 :
                    my_tar.extractall(path="{}/{}".format(self.sdir, "{}_{}".format(self.itmp,i)))
                else:
                    my_tar.extractall(path="{}/{}".format(self.sdir, self.itmp))
                my_tar.close()
            elif magic == '\x50\x4b\x03\x04': # zip
                my_zip = zipfile.ZipFile(data, mode="r")
                if len(inids) > 1 :
                    my_zip.extractall(path="{}/{}".format(self.sdir, "{}_{}".format(self.itmp,i)))
                else:
                    my_zip.extractall(path="{}/{}".format(self.sdir, self.itmp))
                my_zip.close()
            else:
                if len(inids) > 1 :
                    dif = open("{}/{}".format(self.sdir, "{}_{}".format(self.itmp,i)),'w')
                else:
                    dif = open("{}/{}".format(self.sdir, self.itmp),'w')
                dif.write(data.read())
                dif.close()
            rdata.close()
            data.close()

class Validator(TransformBase):
    def __init__(self, args):
        TransformBase.__init__(self,args)
        self.ws_url = args.ws_url
        self.cfg_name = args.cfg_name
        self.sws_id = args.sws_id
        self.etype = args.etype
        self.opt_args = args.opt_args

        # download ws object and find where the validation script is located
        self.wsd = Workspace(url=self.ws_url, token=self.token)
        self.config = self.wsd.get_object({'id' : self.cfg_name, 'workspace' : self.sws_id})['data']['config_map']

        if self.config is None:
            raise Exception("Object {} not found in workspace {}".format(self.cfg_name, self.sws_id))

    def validation_handler (self) :
        ###
        # execute validation
        ## TODO: Add logging

        if self.etype not in self.config['validator']:
          raise Exception("No validation script was registered for {}".format(self.etype))

        fd_list = []
        if os.path.exists("{}/{}".format(self.sdir,self.itmp)):
          fd_list.append( "{}/{}".format(self.sdir,self.itmp))
        else:
          fd_list = glob.glob("{}/{}_*".format(self.sdir,self.itmp))

        for fd in fd_list:
          vcmd_lst = [self.config['validator'][self.etype]['cmd_name'], self.config['validator'][self.etype]['cmd_args']['input'], fd]

          if 'validator' in self.opt_args:
            opt_args = self.opt_args['validator']
            for k in opt_args:
              if k in self.config['validator'][etype]['opt_args'] and opt_args[k] is not None:
                vcmd_lst.append(self.config['validator'][self.etype]['opt_args'][k])
                vcmd_lst.append(opt_args[k])

          p1 = Popen(vcmd_lst, stdout=PIPE)
          out_str = p1.communicate()
          # print output message for error tracking
          if out_str[0] is not None : print out_str[0]
          if out_str[1] is not None : print >> sys.stderr, out_str[1]

          if p1.returncode != 0:
              raise Exception(out_str[1])

class Uploader(Validator):
    def __init__(self, args):
        Validator.__init__(self, args)
        self.kbtype = args.kbtype
        self.ws_id = args.ws_id
        self.outobj_id = args.outobj_id
        self.jid = args.jid
        self.otmp = args.otmp

    def transformation_handler (self) :
        conv_type = "{}-to-{}".format(self.etype, self.kbtype)

        if conv_type not in self.config['transformer']:
          raise Exception("No conversion script was registered for {}".format(conv_type))
        vcmd_lst = [self.config['transformer'][conv_type]['cmd_name']]
        vcmd_lst.append(self.config['transformer'][conv_type]['cmd_args']['input'])
        if 'cmd_args_override' in self.config['transformer'][conv_type] and 'input' in self.config['transformer'][conv_type]['cmd_args_override']:
          if self.config['transformer'][conv_type]['cmd_args_override']['input'] == 'shock_node_id': # use shock node id
            vcmd_lst.append(self.inobj_id)
          else: vcmd_lst.append("{}/{}".format(self.sdir,self.itmp)) # not defined yet
        else: vcmd_lst.append("{}/{}".format(self.sdir,self.itmp)) # default input is the input file or folder

        vcmd_lst.append(self.config['transformer'][conv_type]['cmd_args']['output'])
        if 'cmd_args_override' in self.config['transformer'][conv_type] and 'output' in self.config['transformer'][conv_type]['cmd_args_override']:
          vcmd_lst.append("{}/{}".format(self.sdir,self.otmp)) # not defined yet
        else: vcmd_lst.append("{}/{}".format(self.sdir,self.otmp))

        if 'transformer' in self.opt_args:
          opt_args = self.opt_args['transformer']
          for k in opt_args:
            if k in self.config['transformer'][conv_type]['opt_args'] and opt_args[k] is not None:
              vcmd_lst.append(self.config['transformer'][conv_type]['opt_args'][k])
              vcmd_lst.append(opt_args[k])

        p1 = Popen(vcmd_lst, stdout=PIPE)
        out_str = p1.communicate()
        # print output message for error tracking
        if out_str[0] is not None : print out_str[0]
        if out_str[1] is not None : print >> sys.stderr, out_str[1]

        if p1.returncode != 0:
                raise Exception(out_str[1])

    def upload_handler (self) :

        if self.kbtype in self.config['uploader']: # upload handler is registered
          vcmd_lst = [self.config['uploader'][self.kbtype]['cmd_name'], self.config['uploader'][self.kbtype]['cmd_args']['ws_url'], self.ws_url, self.config['uploader'][self.kbtype]['cmd_args']['ws_id'], self.ws_id, self.config['uploader'][self.kbtype]['cmd_args']['outobj_id'], self.outobj_id,  self.config['uploader'][self.kbtype]['cmd_args']['dir'], self.sdir ]

          if 'uploader' in self.opt_args:
            opt_args = self.opt_args['uploader']
            for k in opt_args:
              if k in self.config['uploader'][self.kbtype]['opt_args'] and opt_args[k] is not None:
                vcmd_lst.append(self.config['uploader'][self.kbtype]['opt_args'][k])
                vcmd_lst.append(opt_args[k])

          print vcmd_lst

          p1 = Popen(vcmd_lst, stdout=PIPE)
          out_str = p1.communicate()
          # print output message for error tracking
          if out_str[0] is not None : print out_str[0]
          if out_str[1] is not None : print >> sys.stderr, out_str[1]

          if p1.returncode != 0:
              raise Exception(out_str[1])
        else: # upload handler was not registered
          self.upload_to_ws() # use default WS uploader

    def upload_to_ws_args (self, sdir,otmp,ws_id,kbtype,outobj_id,inobj_id,etype,jid) :
        jif = open("{}/{}".format(sdir,otmp, 'r'))
        data = json.loads(jif.read())
        jif.close()

        self.wsd.save_objects({'workspace':ws_id, 'objects' : [ {
          'type' : kbtype, 'data' : data, 'name' : outobj_id,
          'meta' : { 'source_id' : inobj_id, 'source_type' : etype,
                     'ujs_job_id' : jid} } ]})

    def upload_to_ws (self) :
        jif = open("{}/{}".format(self.sdir,self.otmp, 'r'))
        data = json.loads(jif.read())
        jif.close()

        self.wsd.save_objects({'workspace':self.ws_id, 'objects' : [ {
          'type' : self.kbtype, 'data' : data, 'name' : self.outobj_id,
          'meta' : { 'source_id' : self.inobj_id, 'source_type' : self.etype,
                     'ujs_job_id' : self.jid} } ]})

class Downloader(TransformBase):
    def __init__(self, args):
        TransformBase.__init__(self,args)
        self.ws_url = args.ws_url
        self.cfg_name = args.cfg_name
        self.sws_id = args.sws_id
        self.etype = args.etype
        self.opt_args = args.opt_args
        self.kbtype = args.kbtype
        #self.otmp = args.otmp
        self.ws_id = args.ws_id
        #self.outobj_id = args.outobj_id
        self.jid = args.jid

        # download ws object and find where the validation script is located
        self.wsd = Workspace(url=self.ws_url, token=self.token)
        self.config = self.wsd.get_object({'id' : self.cfg_name, 'workspace' : self.sws_id})['data']['config_map']

        if self.config is None:
            raise Exception("Object {} not found in workspace {}".format(self.cfg_name, self.sws_id))

    def download_ws_data (self) :
        try:
            os.mkdir(self.sdir)
        except:
            pass

        dif = open("{}/{}".format(self.sdir, "{}".format(self.itmp)),'w')
        data = self.wsd.get_object({'id' : self.inobj_id, 'workspace' : self.ws_id})['data']
        json.dump(data,dif)
        dif.close()

    #def download_handler (ws_url, cfg_name, sws_id, ws_id, in_id, etype, kbtype, sdir, otmp, opt_args, ujs_url, ujs_jid) :
    def download_handler (self) :
        try:
            os.mkdir(self.sdir)
        except:
            pass

        conv_type = "{}-to-{}".format(self.kbtype, self.etype)
        if conv_type  not in self.config['down_transformer'] or 'inobj_id'  not in self.config['down_transformer'][conv_type]['cmd_args'] or 'ws_id'  not in self.config['down_transformer'][conv_type]['cmd_args'] or 'output'  not in self.config['down_transformer'][conv_type]['cmd_args']:
            raise Exception("{} to {} conversion was not properly defined!".format(self.kbtype, self.etype))
        vcmd_lst = [self.config['down_transformer'][conv_type]['cmd_name'],
                    self.config['down_transformer'][conv_type]['cmd_args']['ws_id'], self.ws_id,
                    self.config['down_transformer'][conv_type]['cmd_args']['inobj_id'], self.inobj_id,
                    self.config['down_transformer'][conv_type]['cmd_args']['output'],"{}/{}".format(self.sdir,self.otmp)]

        if 'down_transformer' in self.opt_args:
            opt_args = self.opt_args['down_transformer']
            for k in opt_args:
                if k in self.config['down_transformer'][conv_type]['opt_args'] and opt_args[k] is not None:
                    vcmd_lst.append(self.config['down_transformer'][conv_type]['opt_args'][k])
                    vcmd_lst.append(opt_args[k])

        p1 = Popen(vcmd_lst, stdout=PIPE)
        out_str = p1.communicate()

        if out_str[0] is not None : print out_str[0]
        if out_str[1] is not None : print >> sys.stderr, out_str[1]

        if p1.returncode != 0:
            raise Exception(out_str[1])
<commit_msg>Fixed one more bug for validation<commit_after>#!/usr/bin/python
import argparse
import sys
import os
import os.path
import io
import time
import traceback
import ctypes
import subprocess
from subprocess import Popen, PIPE
import shutil
from optparse import OptionParser
import requests
from requests_toolbelt import MultipartEncoder
import urllib
import urllib2
import json
import tarfile
import zipfile
import glob

try:
    from cStringIO import StringIO
except:
    from StringIO import StringIO

from biokbase.workspace.client import Workspace

BUF_SIZE = 8*1024 # default HTTP LIB client buffer_size

# Base class for Transform service
class TransformBase:

    def __init__(self, args):
        self.shock_url = args.shock_url
        self.inobj_id = args.inobj_id
        self.sdir = args.sdir
        self.itmp = args.itmp
        if(hasattr(args, 'otmp')):
          self.otmp = args.otmp
        else:
          self.otmp = "output"
        self.token = os.environ.get('KB_AUTH_TOKEN')
        self.ssl_verify = True

    def upload_to_shock(self):
        if self.token is None:
            raise Exception("Unable to find token!")
        filePath = "{}/{}".format(self.sdir,self.otmp)

        #build the header
        header = dict()
        header["Authorization"] = "Oauth %s" % self.token

        dataFile = open(os.path.abspath(filePath))
        m = MultipartEncoder(fields={'upload': (os.path.split(filePath)[-1], dataFile)})
        header['Content-Type'] = m.content_type

        try:
            response = requests.post(self.shock_url + "/node", headers=header, data=m, allow_redirects=True, verify=self.ssl_verify)
            dataFile.close()

            if not response.ok:
                response.raise_for_status()

            result = response.json()

            if result['error']:
                raise Exception(result['error'][0])
            else:
                return result["data"]
        except:
            dataFile.close()
            raise

    def download_shock_data(self) :
        if self.token is None:
            raise Exception("Unable to find token!")
        # TODO: Improve folder checking
        try:
            os.mkdir(self.sdir)
        except:
            pass

        inids = self.inobj_id.split(',')

        for i in range(len(inids)):
            meta_req = urllib2.Request("{}/node/{}".format(self.shock_url, inids[i]))
            meta_req.add_header('Authorization',"OAuth {}".format(self.token))
            data_req = urllib2.Request("{}/node/{}?download_raw".format(self.shock_url, inids[i]))
            data_req.add_header('Authorization',"OAuth {}".format(self.token))

            meta = urllib2.urlopen(meta_req)
            md = json.loads(meta.read())
            meta.close()

            rdata = urllib2.urlopen(data_req)

            data = StringIO(rdata.read())
            magic = data.read(4)
            data.seek(0)

            if magic.startswith('\x1f\x8b') or magic.startswith('\x42\x5a') : # gz or bz
                my_tar = tarfile.open(fileobj=data, mode="r|*", bufsize=BUF_SIZE)
                if len(inids) > 1 :
                    my_tar.extractall(path="{}/{}".format(self.sdir, "{}_{}".format(self.itmp,i)))
                else:
                    my_tar.extractall(path="{}/{}".format(self.sdir, self.itmp))
                my_tar.close()
            elif magic == '\x50\x4b\x03\x04': # zip
                my_zip = zipfile.ZipFile(data, mode="r")
                if len(inids) > 1 :
                    my_zip.extractall(path="{}/{}".format(self.sdir, "{}_{}".format(self.itmp,i)))
                else:
                    my_zip.extractall(path="{}/{}".format(self.sdir, self.itmp))
                my_zip.close()
            else:
                if len(inids) > 1 :
                    dif = open("{}/{}".format(self.sdir, "{}_{}".format(self.itmp,i)),'w')
                else:
                    dif = open("{}/{}".format(self.sdir, self.itmp),'w')
                dif.write(data.read())
                dif.close()
            rdata.close()
            data.close()

class Validator(TransformBase):
    def __init__(self, args):
        TransformBase.__init__(self,args)
        self.ws_url = args.ws_url
        self.cfg_name = args.cfg_name
        self.sws_id = args.sws_id
        self.etype = args.etype
        self.opt_args = args.opt_args

        # download ws object and find where the validation script is located
        self.wsd = Workspace(url=self.ws_url, token=self.token)
        self.config = self.wsd.get_object({'id' : self.cfg_name, 'workspace' : self.sws_id})['data']['config_map']

        if self.config is None:
            raise Exception("Object {} not found in workspace {}".format(self.cfg_name, self.sws_id))

    def validation_handler (self) :
        ###
        # execute validation
        ## TODO: Add logging

        if self.etype not in self.config['validator']:
          raise Exception("No validation script was registered for {}".format(self.etype))

        fd_list = []
        if os.path.exists("{}/{}".format(self.sdir,self.itmp)):
          fd_list.append( "{}/{}".format(self.sdir,self.itmp))
        else:
          fd_list = glob.glob("{}/{}_*".format(self.sdir,self.itmp))

        for fd in fd_list:
          vcmd_lst = [self.config['validator'][self.etype]['cmd_name'], self.config['validator'][self.etype]['cmd_args']['input'], fd]

          if 'validator' in self.opt_args:
            opt_args = self.opt_args['validator']
            for k in opt_args:
              if k in self.config['validator'][etype]['opt_args'] and opt_args[k] is not None:
                vcmd_lst.append(self.config['validator'][self.etype]['opt_args'][k])
                vcmd_lst.append(opt_args[k])

          p1 = Popen(vcmd_lst, stdout=PIPE)
          out_str = p1.communicate()
          # print output message for error tracking
          if out_str[0] is not None : print out_str[0]
          if out_str[1] is not None : print >> sys.stderr, out_str[1]

          if p1.returncode != 0:
              raise Exception(out_str[1])

class Uploader(Validator):
    def __init__(self, args):
        Validator.__init__(self, args)
        self.kbtype = args.kbtype
        self.ws_id = args.ws_id
        self.outobj_id = args.outobj_id
        self.jid = args.jid
        self.otmp = args.otmp

    def transformation_handler (self) :
        conv_type = "{}-to-{}".format(self.etype, self.kbtype)

        if conv_type not in self.config['transformer']:
          raise Exception("No conversion script was registered for {}".format(conv_type))
        vcmd_lst = [self.config['transformer'][conv_type]['cmd_name']]
        vcmd_lst.append(self.config['transformer'][conv_type]['cmd_args']['input'])
        if 'cmd_args_override' in self.config['transformer'][conv_type] and 'input' in self.config['transformer'][conv_type]['cmd_args_override']:
          if self.config['transformer'][conv_type]['cmd_args_override']['input'] == 'shock_node_id': # use shock node id
            vcmd_lst.append(self.inobj_id)
          else: vcmd_lst.append("{}/{}".format(self.sdir,self.itmp)) # not defined yet
        else: vcmd_lst.append("{}/{}".format(self.sdir,self.itmp)) # default input is the input file or folder

        vcmd_lst.append(self.config['transformer'][conv_type]['cmd_args']['output'])
        if 'cmd_args_override' in self.config['transformer'][conv_type] and 'output' in self.config['transformer'][conv_type]['cmd_args_override']:
          vcmd_lst.append("{}/{}".format(self.sdir,self.otmp)) # not defined yet
        else: vcmd_lst.append("{}/{}".format(self.sdir,self.otmp))

        if 'transformer' in self.opt_args:
          opt_args = self.opt_args['transformer']
          for k in opt_args:
            if k in self.config['transformer'][conv_type]['opt_args'] and opt_args[k] is not None:
              vcmd_lst.append(self.config['transformer'][conv_type]['opt_args'][k])
              vcmd_lst.append(opt_args[k])

        p1 = Popen(vcmd_lst, stdout=PIPE)
        out_str = p1.communicate()
        # print output message for error tracking
        if out_str[0] is not None : print out_str[0]
        if out_str[1] is not None : print >> sys.stderr, out_str[1]

        if p1.returncode != 0:
                raise Exception(out_str[1])

    def upload_handler (self) :

        if self.kbtype in self.config['uploader']: # upload handler is registered
          vcmd_lst = [self.config['uploader'][self.kbtype]['cmd_name'], self.config['uploader'][self.kbtype]['cmd_args']['ws_url'], self.ws_url, self.config['uploader'][self.kbtype]['cmd_args']['ws_id'], self.ws_id, self.config['uploader'][self.kbtype]['cmd_args']['outobj_id'], self.outobj_id,  self.config['uploader'][self.kbtype]['cmd_args']['dir'], self.sdir ]

          if 'uploader' in self.opt_args:
            opt_args = self.opt_args['uploader']
            for k in opt_args:
              if k in self.config['uploader'][self.kbtype]['opt_args'] and opt_args[k] is not None:
                vcmd_lst.append(self.config['uploader'][self.kbtype]['opt_args'][k])
                vcmd_lst.append(opt_args[k])

          print vcmd_lst

          p1 = Popen(vcmd_lst, stdout=PIPE)
          out_str = p1.communicate()
          # print output message for error tracking
          if out_str[0] is not None : print out_str[0]
          if out_str[1] is not None : print >> sys.stderr, out_str[1]

          if p1.returncode != 0:
              raise Exception(out_str[1])
        else: # upload handler was not registered
          self.upload_to_ws() # use default WS uploader

    def upload_to_ws_args (self, sdir,otmp,ws_id,kbtype,outobj_id,inobj_id,etype,jid) :
        jif = open("{}/{}".format(sdir,otmp, 'r'))
        data = json.loads(jif.read())
        jif.close()

        self.wsd.save_objects({'workspace':ws_id, 'objects' : [ {
          'type' : kbtype, 'data' : data, 'name' : outobj_id,
          'meta' : { 'source_id' : inobj_id, 'source_type' : etype,
                     'ujs_job_id' : jid} } ]})

    def upload_to_ws (self) :
        jif = open("{}/{}".format(self.sdir,self.otmp, 'r'))
        data = json.loads(jif.read())
        jif.close()

        self.wsd.save_objects({'workspace':self.ws_id, 'objects' : [ {
          'type' : self.kbtype, 'data' : data, 'name' : self.outobj_id,
          'meta' : { 'source_id' : self.inobj_id, 'source_type' : self.etype,
                     'ujs_job_id' : self.jid} } ]})

class Downloader(TransformBase):
    def __init__(self, args):
        TransformBase.__init__(self,args)
        self.ws_url = args.ws_url
        self.cfg_name = args.cfg_name
        self.sws_id = args.sws_id
        self.etype = args.etype
        self.opt_args = args.opt_args
        self.kbtype = args.kbtype
        #self.otmp = args.otmp
        self.ws_id = args.ws_id
        #self.outobj_id = args.outobj_id
        self.jid = args.jid

        # download ws object and find where the validation script is located
        self.wsd = Workspace(url=self.ws_url, token=self.token)
        self.config = self.wsd.get_object({'id' : self.cfg_name, 'workspace' : self.sws_id})['data']['config_map']

        if self.config is None:
            raise Exception("Object {} not found in workspace {}".format(self.cfg_name, self.sws_id))

    def download_ws_data (self) :
        try:
            os.mkdir(self.sdir)
        except:
            pass

        dif = open("{}/{}".format(self.sdir, "{}".format(self.itmp)),'w')
        data = self.wsd.get_object({'id' : self.inobj_id, 'workspace' : self.ws_id})['data']
        json.dump(data,dif)
        dif.close()

    #def download_handler (ws_url, cfg_name, sws_id, ws_id, in_id, etype, kbtype, sdir, otmp, opt_args, ujs_url, ujs_jid) :
    def download_handler (self) :
        try:
            os.mkdir(self.sdir)
        except:
            pass

        conv_type = "{}-to-{}".format(self.kbtype, self.etype)
        if conv_type  not in self.config['down_transformer'] or 'inobj_id'  not in self.config['down_transformer'][conv_type]['cmd_args'] or 'ws_id'  not in self.config['down_transformer'][conv_type]['cmd_args'] or 'output'  not in self.config['down_transformer'][conv_type]['cmd_args']:
            raise Exception("{} to {} conversion was not properly defined!".format(self.kbtype, self.etype))
        vcmd_lst = [self.config['down_transformer'][conv_type]['cmd_name'],
                    self.config['down_transformer'][conv_type]['cmd_args']['ws_id'], self.ws_id,
                    self.config['down_transformer'][conv_type]['cmd_args']['inobj_id'], self.inobj_id,
                    self.config['down_transformer'][conv_type]['cmd_args']['output'],"{}/{}".format(self.sdir,self.otmp)]

        if 'down_transformer' in self.opt_args:
            opt_args = self.opt_args['down_transformer']
            for k in opt_args:
                if k in self.config['down_transformer'][conv_type]['opt_args'] and opt_args[k] is not None:
                    vcmd_lst.append(self.config['down_transformer'][conv_type]['opt_args'][k])
                    vcmd_lst.append(opt_args[k])

        p1 = Popen(vcmd_lst, stdout=PIPE)
        out_str = p1.communicate()

        if out_str[0] is not None : print out_str[0]
        if out_str[1] is not None : print >> sys.stderr, out_str[1]

        if p1.returncode != 0:
            raise Exception(out_str[1])
<|endoftext|>


<commit_before>import json
import logging
from collections import OrderedDict
from pprint import pformat

import click

from clifunzone import click_utils
from clifunzone import xml_utils
from clifunzone.reflection_utils import varsdict

try:
    from lxml import etree as ET
except ImportError:
    try:
        import xml.etree.cElementTree as ET
    except ImportError:
        import xml.etree.ElementTree as ET
# import xml.etree.ElementTree as ET

def debug_context(**kwargs):
    ctx = click.get_current_context()

    # debug_ = ctx.params.get('debug') or ctx.parent.params.get('debug') if ctx.parent else False
    # if debug_:
    #     echo_context(ctx)

    click_utils.inherit_parent_params(ctx, ('debug',))

    debug_ = ctx.params['debug']

    if debug_:
        click.echo('Debug mode: %s' % ('enabled' if debug_ else 'disabled'))
        click_utils.echo_context(ctx)
        click_utils.echo_kwargs(kwargs)

        subcommand = ctx.invoked_subcommand
        click.echo('subcommand: %s' % subcommand)

@click.group(context_settings=click_utils.CONTEXT_SETTINGS, invoke_without_command=True)
@click.version_option(version='1.0.0')
@click.option('--debug/--silent', '-d/-s', 'debug', default=False)
# @click.option('--debug', '-d', 'debug', flag_value=True, default=True)
# @click.option('--silent', '-s', 'debug', flag_value=False)
def cli(debug):
    """
    Provides CLI commands for interacting with Robot Framework test output data/files (XML format).
    """
    ctx = click.get_current_context()
    if debug:
        click_utils.echo_context(ctx)

    subcommand = ctx.invoked_subcommand
    # click.echo('Subcommand: {}'.format(subcommand))
    if subcommand is None:
        click.echo('I was invoked without a subcommand...')
    else:
        if debug:
            click.echo('I was invoked with subcommand: %s.' % subcommand)

    try:
        debug_context(**{})
    except:
        logging.exception('debug_context error')

@cli.command()
@click.option('--input', '-i', type=click.Path(exists=True, dir_okay=False, allow_dash=True),
              help="the path to the file containing the input. Or '-' to use stdin (e.g. piped input).")
@click.option('--verbose', '-v', is_flag=True, type=click.BOOL,
              help='enables more detailed output.')
def info(input, verbose, **kwargs):
    """
    Provides info about the input. Requires valid input.
    """
    if not input:
        input = '-'
    with click.open_file(input, mode='rb') as f:
        # root = xml_utils.load(f)
        tree = ET.parse(f)
        root = tree.getroot()
        d = {}
        d.update({'xml': xml_utils.element_info(root)})

        rf_metrics = {
            'suites': xml_utils.count_elements(tree, xpath='//suite'),
            'tests': xml_utils.count_elements(tree, xpath='//test'),
            'messages': xml_utils.count_elements(tree, xpath='//msg')
        }
        d.update({'robot': rf_metrics})

        if verbose:
            d['_object'] = {
                'type': type(root),
                # 'repr': repr(root),
                # 'vars': sorted(vars(root)),
                # 'dir': sorted(dir(root)),
                'members': sorted(varsdict(root).keys())
            }
        # click.echo(d)
        # click.echo(sorted(d.items()))
        if verbose:
            s = pformat(d)
        else:
            s = json.dumps(d, indent=2, sort_keys=True)
        click.echo(s)

@cli.command(short_help='echo the unparsed input')
@click.option('--input', '-i', type=click.Path(exists=True, dir_okay=False, allow_dash=True),
              help="the path to the file containing the input to be echoed. Or '-' to use stdin (e.g. piped input).")
def echo(input, **kwargs):
    """
    Echo the input.
    """
    if not input:
        input = '-'
    with click.open_file(input, mode='rb') as f:
        s = f.read()
        click.echo(s)

@cli.command()
@click.option('--input', '-i', type=click.Path(exists=True, dir_okay=False, allow_dash=True),
              help="the path to the file containing the input. Or '-' to use stdin (e.g. piped input).")
@click.option('--start', type=click.INT)
@click.option('--stop', type=click.INT)
@click.option('--step', type=click.INT)
@click.option('--status', type=click.Choice(['PASS', 'FAIL']))
@click.option('--showargs', '-a', is_flag=True, default=False,
              help="display each test's args.")
@click.option('--showmsgs', '-m', is_flag=True, default=False,
              help="display each test's inner messages.")
@click.option('--pretty', '-p', is_flag=True, default=False, help='pretty format')
def keywords(input, start, stop, step, status, showargs, showmsgs, pretty, **kwargs):
    """
    Extracts RF keywords from the input.
    """
    if not input:
        input = '-'
    with click.open_file(input, mode='rb') as f:
        tree = ET.parse(f)
        root = tree.getroot()
        kws = [elem for elem in root.iter(tag='kw')]
        kws = [node_to_dict(kw, showargs=showargs, showmsgs=showmsgs) for kw in kws]
        if status:
            kws = [kw for kw in kws if kw['status']['status'] == status]
        kws = kws[start:stop:step]
        if pretty:
            output = json.dumps(kws, indent=4)
            click.echo(output)
        else:
            # click.echo(json.dumps(kws))
            for kw in kws:
                kw = json.dumps(kw)
                click.echo(kw)

@cli.command(name='tests')
@click.option('--input', '-i', type=click.Path(exists=True, dir_okay=False, allow_dash=True),
              help="the path to the file containing the input. Or '-' to use stdin (e.g. piped input).")
@click.option('--start', type=click.INT)
@click.option('--stop', type=click.INT)
@click.option('--step', type=click.INT)
@click.option('--status', type=click.Choice(['PASS', 'FAIL']))
@click.option('--showargs', '-a', is_flag=True, default=False,
              help="display each test's args.")
@click.option('--showmsgs', '-m', is_flag=True, default=False,
              help="display each test's inner messages.")
@click.option('--pretty', '-p', is_flag=True, default=False, help='pretty format')
def testscommand(input, start, stop, step, status, showargs, showmsgs, pretty, **kwargs):
    """
    Extracts RF test data from the input.
    """
    if not input:
        input = '-'
    with click.open_file(input, mode='rb') as f:
        tree = ET.parse(f)
        root = tree.getroot()
        test_elements = [elem for elem in root.iter(tag='test')]
        test_elements = [node_to_dict(test_element, showargs=showargs, showmsgs=showmsgs)
                         for test_element in test_elements]
        if status:
            test_elements = [test_element for test_element in test_elements if
                             test_element['status']['status'] == status]
        test_elements = test_elements[start:stop:step]
        if pretty:
            output = json.dumps(test_elements, indent=4)
            click.echo(output)
        else:
            # click.echo(json.dumps(kws))
            for test_element in test_elements:
                test_element = json.dumps(test_element)
                click.echo(test_element)

def node_to_dict(root, showargs=False, showmsgs=False, showtags=True):
    d = OrderedDict()

    # attribs = [k for k in ['id', 'name', 'type'] if k in root.attrib]
    attribs = root.attrib.keys()
    for attrib in attribs:
        d['@' + attrib] = root.attrib[attrib]

    status = root.find('status').attrib
    if status is not None:
        if type(status) is not dict:
            # Note: if lxml is being used, status will be a non-JSON-serializable object (<type 'lxml.etree._Attrib'>).
            # Therefore, we must manually convert it to a regular dict to enable JSON serialization compatibility.
            status = {k: v for k, v in status.items()}
        d['status'] = status

    doc = root.find('doc')
    if doc is not None:
        d['doc'] = doc.text

    if showtags:
        tags = root.find('tags')
        if tags is not None:
            tags = tags.iter(tag='tag')
            tags = [tag.text for tag in tags]
        d['tags'] = tags

    if showargs:
        args = root.find('arguments')
        if args is not None:
            args = args.iter(tag='arg')
            args = [arg.text for arg in args]
        d['args'] = args

    if showmsgs:
        msgs = root.iter(tag='msg')
        if msgs is not None:
            msgs = [msg.text for msg in msgs]
        d['messages'] = msgs

    return d

def main():
    cli()

if __name__ == '__main__':
    # main(default_map={
    #     'info': {
    #         'input': 'rfxml_parse.input.01.xml'
    #     }
    # })
    main()
<commit_msg>rfxmltool: echo command: updated command-level help text.<commit_after>import json
import logging
from collections import OrderedDict
from pprint import pformat

import click

from clifunzone import click_utils
from clifunzone import xml_utils
from clifunzone.reflection_utils import varsdict

try:
    from lxml import etree as ET
except ImportError:
    try:
        import xml.etree.cElementTree as ET
    except ImportError:
        import xml.etree.ElementTree as ET
# import xml.etree.ElementTree as ET

def debug_context(**kwargs):
    ctx = click.get_current_context()

    # debug_ = ctx.params.get('debug') or ctx.parent.params.get('debug') if ctx.parent else False
    # if debug_:
    #     echo_context(ctx)

    click_utils.inherit_parent_params(ctx, ('debug',))

    debug_ = ctx.params['debug']

    if debug_:
        click.echo('Debug mode: %s' % ('enabled' if debug_ else 'disabled'))
        click_utils.echo_context(ctx)
        click_utils.echo_kwargs(kwargs)

        subcommand = ctx.invoked_subcommand
        click.echo('subcommand: %s' % subcommand)

@click.group(context_settings=click_utils.CONTEXT_SETTINGS, invoke_without_command=True)
@click.version_option(version='1.0.0')
@click.option('--debug/--silent', '-d/-s', 'debug', default=False)
# @click.option('--debug', '-d', 'debug', flag_value=True, default=True)
# @click.option('--silent', '-s', 'debug', flag_value=False)
def cli(debug):
    """
    Provides CLI commands for interacting with Robot Framework test output data/files (XML format).
    """
    ctx = click.get_current_context()
    if debug:
        click_utils.echo_context(ctx)

    subcommand = ctx.invoked_subcommand
    # click.echo('Subcommand: {}'.format(subcommand))
    if subcommand is None:
        click.echo('I was invoked without a subcommand...')
    else:
        if debug:
            click.echo('I was invoked with subcommand: %s.' % subcommand)

    try:
        debug_context(**{})
    except:
        logging.exception('debug_context error')

@cli.command()
@click.option('--input', '-i', type=click.Path(exists=True, dir_okay=False, allow_dash=True),
              help="the path to the file containing the input. Or '-' to use stdin (e.g. piped input).")
@click.option('--verbose', '-v', is_flag=True, type=click.BOOL,
              help='enables more detailed output.')
def info(input, verbose, **kwargs):
    """
    Provides info about the input. Requires valid input.
    """
    if not input:
        input = '-'
    with click.open_file(input, mode='rb') as f:
        # root = xml_utils.load(f)
        tree = ET.parse(f)
        root = tree.getroot()
        d = {}
        d.update({'xml': xml_utils.element_info(root)})

        rf_metrics = {
            'suites': xml_utils.count_elements(tree, xpath='//suite'),
            'tests': xml_utils.count_elements(tree, xpath='//test'),
            'messages': xml_utils.count_elements(tree, xpath='//msg')
        }
        d.update({'robot': rf_metrics})

        if verbose:
            d['_object'] = {
                'type': type(root),
                # 'repr': repr(root),
                # 'vars': sorted(vars(root)),
                # 'dir': sorted(dir(root)),
                'members': sorted(varsdict(root).keys())
            }
        # click.echo(d)
        # click.echo(sorted(d.items()))
        if verbose:
            s = pformat(d)
        else:
            s = json.dumps(d, indent=2, sort_keys=True)
        click.echo(s)

@cli.command(short_help='echo the unparsed input')
@click.option('--input', '-i', type=click.Path(exists=True, dir_okay=False, allow_dash=True),
              help="the path to the file containing the input. Or '-' to use stdin (e.g. piped input).")
def echo(input, **kwargs):
    """
    Echo the (unparsed) input.
    """
    if not input:
        input = '-'
    with click.open_file(input, mode='rb') as f:
        s = f.read()
        click.echo(s)

@cli.command()
@click.option('--input', '-i', type=click.Path(exists=True, dir_okay=False, allow_dash=True),
              help="the path to the file containing the input. Or '-' to use stdin (e.g. piped input).")
@click.option('--start', type=click.INT)
@click.option('--stop', type=click.INT)
@click.option('--step', type=click.INT)
@click.option('--status', type=click.Choice(['PASS', 'FAIL']))
@click.option('--showargs', '-a', is_flag=True, default=False,
              help="display each test's args.")
@click.option('--showmsgs', '-m', is_flag=True, default=False,
              help="display each test's inner messages.")
@click.option('--pretty', '-p', is_flag=True, default=False, help='pretty format')
def keywords(input, start, stop, step, status, showargs, showmsgs, pretty, **kwargs):
    """
    Extracts RF keywords from the input.
    """
    if not input:
        input = '-'
    with click.open_file(input, mode='rb') as f:
        tree = ET.parse(f)
        root = tree.getroot()
        kws = [elem for elem in root.iter(tag='kw')]
        kws = [node_to_dict(kw, showargs=showargs, showmsgs=showmsgs) for kw in kws]
        if status:
            kws = [kw for kw in kws if kw['status']['status'] == status]
        kws = kws[start:stop:step]
        if pretty:
            output = json.dumps(kws, indent=4)
            click.echo(output)
        else:
            # click.echo(json.dumps(kws))
            for kw in kws:
                kw = json.dumps(kw)
                click.echo(kw)

@cli.command(name='tests')
@click.option('--input', '-i', type=click.Path(exists=True, dir_okay=False, allow_dash=True),
              help="the path to the file containing the input. Or '-' to use stdin (e.g. piped input).")
@click.option('--start', type=click.INT)
@click.option('--stop', type=click.INT)
@click.option('--step', type=click.INT)
@click.option('--status', type=click.Choice(['PASS', 'FAIL']))
@click.option('--showargs', '-a', is_flag=True, default=False,
              help="display each test's args.")
@click.option('--showmsgs', '-m', is_flag=True, default=False,
              help="display each test's inner messages.")
@click.option('--pretty', '-p', is_flag=True, default=False, help='pretty format')
def testscommand(input, start, stop, step, status, showargs, showmsgs, pretty, **kwargs):
    """
    Extracts RF test data from the input.
    """
    if not input:
        input = '-'
    with click.open_file(input, mode='rb') as f:
        tree = ET.parse(f)
        root = tree.getroot()
        test_elements = [elem for elem in root.iter(tag='test')]
        test_elements = [node_to_dict(test_element, showargs=showargs, showmsgs=showmsgs)
                         for test_element in test_elements]
        if status:
            test_elements = [test_element for test_element in test_elements if
                             test_element['status']['status'] == status]
        test_elements = test_elements[start:stop:step]
        if pretty:
            output = json.dumps(test_elements, indent=4)
            click.echo(output)
        else:
            # click.echo(json.dumps(kws))
            for test_element in test_elements:
                test_element = json.dumps(test_element)
                click.echo(test_element)

def node_to_dict(root, showargs=False, showmsgs=False, showtags=True):
    d = OrderedDict()

    # attribs = [k for k in ['id', 'name', 'type'] if k in root.attrib]
    attribs = root.attrib.keys()
    for attrib in attribs:
        d['@' + attrib] = root.attrib[attrib]

    status = root.find('status').attrib
    if status is not None:
        if type(status) is not dict:
            # Note: if lxml is being used, status will be a non-JSON-serializable object (<type 'lxml.etree._Attrib'>).
            # Therefore, we must manually convert it to a regular dict to enable JSON serialization compatibility.
            status = {k: v for k, v in status.items()}
        d['status'] = status

    doc = root.find('doc')
    if doc is not None:
        d['doc'] = doc.text

    if showtags:
        tags = root.find('tags')
        if tags is not None:
            tags = tags.iter(tag='tag')
            tags = [tag.text for tag in tags]
        d['tags'] = tags

    if showargs:
        args = root.find('arguments')
        if args is not None:
            args = args.iter(tag='arg')
            args = [arg.text for arg in args]
        d['args'] = args

    if showmsgs:
        msgs = root.iter(tag='msg')
        if msgs is not None:
            msgs = [msg.text for msg in msgs]
        d['messages'] = msgs

    return d

def main():
    cli()

if __name__ == '__main__':
    # main(default_map={
    #     'info': {
    #         'input': 'rfxml_parse.input.01.xml'
    #     }
    # })
    main()
<|endoftext|>