<commit_before>from zeit.cms.content.interfaces import WRITEABLE_ALWAYS
from zeit.retresco.tag import Tag
import grokcore.component as grok
import logging
import lxml.etree
import lxml.objectify
import zeit.cms.content.dav
import zeit.cms.tagging.interfaces
import zeit.connector.interfaces
import zeit.retresco.interfaces
import zope.component

NAMESPACE = "http://namespaces.zeit.de/CMS/tagging"
KEYWORD_PROPERTY = ('rankedTags', NAMESPACE)
DISABLED_PROPERTY = ('disabled', NAMESPACE)
DISABLED_SEPARATOR = '\x09'
PINNED_PROPERTY = ('pinned', NAMESPACE)

# To speed up indexing we do not checkout resources to update properties.
# This is why we make the keyword property writeable.
property_manager = zeit.cms.content.liveproperty.LiveProperties
property_manager.register_live_property(
    KEYWORD_PROPERTY[0], NAMESPACE, WRITEABLE_ALWAYS)
property_manager.register_live_property(
    DISABLED_PROPERTY[0], NAMESPACE, WRITEABLE_ALWAYS)
property_manager.register_live_property(
    KEYWORD_PROPERTY[0], NAMESPACE, WRITEABLE_ALWAYS)
property_manager.register_live_property(
    PINNED_PROPERTY[0], NAMESPACE, WRITEABLE_ALWAYS)

log = logging.getLogger(__name__)

class Tagger(zeit.cms.content.dav.DAVPropertiesAdapter):
    """Serializes keywords to an XML structure and stores it in a DAV-Property.

    The serialization format (including pinned and disabled keywords) is
    actually backwards compatible with zeit.intrafind, but since the TMS uses
    different IDs for keywords, switching to retresco means that
    pinned+disabled keywords will effectively be reset. (Keywords that are
    assigned to content already will remain usable, though, since the label is
    pretty much the only interesting property.)

    We need to read/write all tags from/to xml for every operation, since
    adapting to ``ITagger`` returns a *different* Tagger each time, so we
    cannot (easily) cache anything on the Tagger instance.

    There are also '<rankedTags>' in the XML of ``context`` but these are
    written by ``zeit.cms.tagging.tag.add_ranked_tags_to_head``. At this point,
    the information about pinned and disabled tags is omitted, as it is only
    needed in vivi.

    """

    grok.implements(zeit.cms.tagging.interfaces.ITagger)

    def __iter__(self):
        tags = self.to_xml()
        if tags is None:
            return iter(())
        return (Tag(x.text, x.get('type', '')).code
                for x in tags.iterchildren())

    def __len__(self):
        return len(list(self.__iter__()))

    def __getitem__(self, key):
        node = self._find_tag_node(key)
        tag = self._create_tag(unicode(node), node.get('type', ''))
        return tag

    def __setitem__(self, key, value):
        tags = self.to_xml()
        if tags is None:
            E = lxml.objectify.ElementMaker(namespace=NAMESPACE)
            root = E.rankedTags()
            tags = E.rankedTags()
            root.append(tags)
        # XXX the handling of namespaces here seems chaotic
        E = lxml.objectify.ElementMaker()
        tags.append(E.tag(value.label, type=value.entity_type or ''))
        dav = zeit.connector.interfaces.IWebDAVProperties(self)
        dav[KEYWORD_PROPERTY] = lxml.etree.tostring(tags.getroottree())

    def values(self):
        tags = self.to_xml()
        if tags is None:
            return iter(())
        return (self._create_tag(unicode(node), node.get('type', ''))
                for node in tags.iterchildren())

    def get(self, key, default=None):
        try:
            return self[key]
        except KeyError:
            return default

    def keys(self):
        raise NotImplementedError()

    def items(self):
        raise NotImplementedError()

    def __contains__(self, key):
        return key in list(self)

    def updateOrder(self, keys):
        keys = list(keys)  # in case we've got a generator
        if set(keys) != set(self):
            raise ValueError(
                'Must pass in the same keys already present %r, not %r'
                % (list(self), keys))
        orderd = []
        tags = self.to_xml()
        for key in keys:
            tag = self._find_tag_node(key, tags)
            orderd.append(tag)
            tags.remove(tag)
        tags.extend(orderd)
        dav = zeit.connector.interfaces.IWebDAVProperties(self)
        dav[KEYWORD_PROPERTY] = lxml.etree.tostring(tags.getroottree())

    def __delitem__(self, key):
        node = self._find_tag_node(key)
        node.getparent().remove(node)
        dav = zeit.connector.interfaces.IWebDAVProperties(self)
        dav[KEYWORD_PROPERTY] = lxml.etree.tostring(node.getroottree())
        self._disable(key)

    def set_pinned(self, keys):
        dav = zeit.connector.interfaces.IWebDAVProperties(self)
        dav[PINNED_PROPERTY] = DISABLED_SEPARATOR.join(keys)

    @property
    def pinned(self):
        dav = zeit.connector.interfaces.IWebDAVProperties(self)
        value = dav.get(PINNED_PROPERTY, '')
        if not value:
            return ()
        return tuple(value.split(DISABLED_SEPARATOR))

    @property
    def disabled(self):
        dav = zeit.connector.interfaces.IWebDAVProperties(self)
        disabled = dav.get(DISABLED_PROPERTY)
        if not disabled:
            return ()
        return tuple(disabled.split(DISABLED_SEPARATOR))

    def _disable(self, key):
        dav = zeit.connector.interfaces.IWebDAVProperties(self)
        disabled = self.disabled + (key,)
        dav[DISABLED_PROPERTY] = DISABLED_SEPARATOR.join(disabled)

    def to_xml(self):
        dav = zeit.connector.interfaces.IWebDAVProperties(self)
        try:
            tags = lxml.objectify.fromstring(dav.get(KEYWORD_PROPERTY, ''))
        except lxml.etree.XMLSyntaxError:
            return None
        if tags.tag != '{{{1}}}{0}'.format(*KEYWORD_PROPERTY):
            return None
        return tags.getchildren()[0]

    def _find_tag_node(self, key, tags=None):
        if tags is None:
            tags = self.to_xml()
        if tags is None:
            raise KeyError(key)

        node = [x for x in tags.iterchildren()
                if Tag(x.text, x.get('type', '')).code == key]
        if not node:
            node = [x for x in tags.iterchildren() if x.get('uuid') == key]
        if not node:
            raise KeyError(key)
        return node[0]

    def _create_tag(self, label, entity_type):
        tag = Tag(label, entity_type)
        if tag.code in self.pinned:
            tag.pinned = True
        # XXX it is not understood by the writer, whether we really need this.
        # It was kept from the former implementation in `zeit.cms.tagging`.
        tag.__parent__ = self
        return tag

    def _map_pinned_codes(self):
        mapping = zeit.retresco.convert.CommonMetadata.entity_types
        xml = self.to_xml()
        pinned_map = {}

        for pin in self.pinned:
            tag = xml.find("tag[@uuid='{}']".format(pin))
            if tag:
                pinned_map[tag.get('uuid')] = (
                    Tag(tag.text, mapping.get(tag.get('type', ''), '')).code)
            else:
                pinned_map[pin] = pin
        return pinned_map

    def update(self, update_with=None):
        """Update the keywords with generated keywords from retresco.

        A number of reasonable keywords are retrieved from retresco. This set
        is reduced by the disabled keywords and enriched by the pinned
        keywords. The resulting set is finally written to the DAV property.

        """

        log.info('Updating tags for %s', self.context.uniqueId)
        tms = zope.component.getUtility(zeit.retresco.interfaces.ITMS)

        if update_with is not None:
            keywords = tms.generate_keyword_list(update_with)
        else:
            keywords = tms.extract_keywords(self.context)

        E = lxml.objectify.ElementMaker(namespace=NAMESPACE)
        root = E.rankedTags()
        new_tags = E.rankedTags()
        root.append(new_tags)

        # XXX the handling of namespaces here seems chaotic
        E = lxml.objectify.ElementMaker()
        new_codes = set()

        for tag in keywords:
            if tag.code in self.disabled:
                continue
            new_codes.add(tag.code)
            new_tags.append(E.tag(tag.label, type=tag.entity_type or ''))

        old_tags = self.to_xml()
        mapped_pinned_codes = self._map_pinned_codes()
        new_pinned_codes = []
        for old_code, code in mapped_pinned_codes.iteritems():
            if code not in new_codes:
                try:
                    tag = self._find_tag_node(old_code, old_tags)
                    mapping = zeit.retresco.convert.CommonMetadata.entity_types
                    tag_type = mapping.get(tag.get('type', ''), '')
                    new_tags.append(E.tag(tag.text, type=tag_type))
                except KeyError:
                    pass
            new_pinned_codes.append(code)

        dav = zeit.connector.interfaces.IWebDAVProperties(self)
        dav[KEYWORD_PROPERTY] = lxml.etree.tostring(root.getroottree())
        dav[PINNED_PROPERTY] = "\t".join(new_pinned_codes)
        dav[DISABLED_PROPERTY] = u''
<commit_msg>Fix test<commit_after>from zeit.cms.content.interfaces import WRITEABLE_ALWAYS
from zeit.retresco.tag import Tag
import grokcore.component as grok
import logging
import lxml.etree
import lxml.objectify
import zeit.cms.content.dav
import zeit.cms.tagging.interfaces
import zeit.connector.interfaces
import zeit.retresco.interfaces
import zope.component

NAMESPACE = "http://namespaces.zeit.de/CMS/tagging"
KEYWORD_PROPERTY = ('rankedTags', NAMESPACE)
DISABLED_PROPERTY = ('disabled', NAMESPACE)
DISABLED_SEPARATOR = '\x09'
PINNED_PROPERTY = ('pinned', NAMESPACE)

# To speed up indexing we do not checkout resources to update properties.
# This is why we make the keyword property writeable.
property_manager = zeit.cms.content.liveproperty.LiveProperties
property_manager.register_live_property(
    KEYWORD_PROPERTY[0], NAMESPACE, WRITEABLE_ALWAYS)
property_manager.register_live_property(
    DISABLED_PROPERTY[0], NAMESPACE, WRITEABLE_ALWAYS)
property_manager.register_live_property(
    KEYWORD_PROPERTY[0], NAMESPACE, WRITEABLE_ALWAYS)
property_manager.register_live_property(
    PINNED_PROPERTY[0], NAMESPACE, WRITEABLE_ALWAYS)

log = logging.getLogger(__name__)

class Tagger(zeit.cms.content.dav.DAVPropertiesAdapter):
    """Serializes keywords to an XML structure and stores it in a DAV-Property.

    The serialization format (including pinned and disabled keywords) is
    actually backwards compatible with zeit.intrafind, but since the TMS uses
    different IDs for keywords, switching to retresco means that
    pinned+disabled keywords will effectively be reset. (Keywords that are
    assigned to content already will remain usable, though, since the label is
    pretty much the only interesting property.)

    We need to read/write all tags from/to xml for every operation, since
    adapting to ``ITagger`` returns a *different* Tagger each time, so we
    cannot (easily) cache anything on the Tagger instance.

    There are also '<rankedTags>' in the XML of ``context`` but these are
    written by ``zeit.cms.tagging.tag.add_ranked_tags_to_head``. At this point,
    the information about pinned and disabled tags is omitted, as it is only
    needed in vivi.

    """

    grok.implements(zeit.cms.tagging.interfaces.ITagger)

    def __iter__(self):
        tags = self.to_xml()
        if tags is None:
            return iter(())
        return (Tag(x.text, x.get('type', '')).code
                for x in tags.iterchildren())

    def __len__(self):
        return len(list(self.__iter__()))

    def __getitem__(self, key):
        node = self._find_tag_node(key)
        tag = self._create_tag(unicode(node), node.get('type', ''))
        return tag

    def __setitem__(self, key, value):
        tags = self.to_xml()
        if tags is None:
            E = lxml.objectify.ElementMaker(namespace=NAMESPACE)
            root = E.rankedTags()
            tags = E.rankedTags()
            root.append(tags)
        # XXX the handling of namespaces here seems chaotic
        E = lxml.objectify.ElementMaker()
        tags.append(E.tag(value.label, type=value.entity_type or ''))
        dav = zeit.connector.interfaces.IWebDAVProperties(self)
        dav[KEYWORD_PROPERTY] = lxml.etree.tostring(tags.getroottree())

    def values(self):
        tags = self.to_xml()
        if tags is None:
            return iter(())
        return (self._create_tag(unicode(node), node.get('type', ''))
                for node in tags.iterchildren())

    def get(self, key, default=None):
        try:
            return self[key]
        except KeyError:
            return default

    def keys(self):
        raise NotImplementedError()

    def items(self):
        raise NotImplementedError()

    def __contains__(self, key):
        return key in list(self)

    def updateOrder(self, keys):
        keys = list(keys)  # in case we've got a generator
        if set(keys) != set(self):
            raise ValueError(
                'Must pass in the same keys already present %r, not %r'
                % (list(self), keys))
        orderd = []
        tags = self.to_xml()
        for key in keys:
            tag = self._find_tag_node(key, tags)
            orderd.append(tag)
            tags.remove(tag)
        tags.extend(orderd)
        dav = zeit.connector.interfaces.IWebDAVProperties(self)
        dav[KEYWORD_PROPERTY] = lxml.etree.tostring(tags.getroottree())

    def __delitem__(self, key):
        node = self._find_tag_node(key)
        node.getparent().remove(node)
        dav = zeit.connector.interfaces.IWebDAVProperties(self)
        dav[KEYWORD_PROPERTY] = lxml.etree.tostring(node.getroottree())
        self._disable(key)

    def set_pinned(self, keys):
        dav = zeit.connector.interfaces.IWebDAVProperties(self)
        dav[PINNED_PROPERTY] = DISABLED_SEPARATOR.join(keys)

    @property
    def pinned(self):
        dav = zeit.connector.interfaces.IWebDAVProperties(self)
        value = dav.get(PINNED_PROPERTY, '')
        if not value:
            return ()
        return tuple(value.split(DISABLED_SEPARATOR))

    @property
    def disabled(self):
        dav = zeit.connector.interfaces.IWebDAVProperties(self)
        disabled = dav.get(DISABLED_PROPERTY)
        if not disabled:
            return ()
        return tuple(disabled.split(DISABLED_SEPARATOR))

    def _disable(self, key):
        dav = zeit.connector.interfaces.IWebDAVProperties(self)
        disabled = self.disabled + (key,)
        dav[DISABLED_PROPERTY] = DISABLED_SEPARATOR.join(disabled)

    def to_xml(self):
        dav = zeit.connector.interfaces.IWebDAVProperties(self)
        try:
            tags = lxml.objectify.fromstring(dav.get(KEYWORD_PROPERTY, ''))
        except lxml.etree.XMLSyntaxError:
            return None
        if tags.tag != '{{{1}}}{0}'.format(*KEYWORD_PROPERTY):
            return None
        return tags.getchildren()[0]

    def _find_tag_node(self, key, tags=None):
        if tags is None:
            tags = self.to_xml()
        if tags is None:
            raise KeyError(key)

        node = [x for x in tags.iterchildren()
                if Tag(x.text, x.get('type', '')).code == key]
        if not node:
            node = [x for x in tags.iterchildren() if x.get('uuid') == key]
        if not node:
            raise KeyError(key)
        return node[0]

    def _create_tag(self, label, entity_type):
        tag = Tag(label, entity_type)
        if tag.code in self.pinned:
            tag.pinned = True
        # XXX it is not understood by the writer, whether we really need this.
        # It was kept from the former implementation in `zeit.cms.tagging`.
        tag.__parent__ = self
        return tag

    def _map_pinned_codes(self):
        mapping = zeit.retresco.convert.CommonMetadata.entity_types
        xml = self.to_xml()
        pinned_map = {}

        for pin in self.pinned:
            tag = xml.find(u"tag[@uuid='{}']".format(pin))
            if tag:
                pinned_map[tag.get('uuid')] = (
                    Tag(tag.text, mapping.get(tag.get('type', ''), '')).code)
            else:
                pinned_map[pin] = pin
        return pinned_map

    def update(self, update_with=None):
        """Update the keywords with generated keywords from retresco.

        A number of reasonable keywords are retrieved from retresco. This set
        is reduced by the disabled keywords and enriched by the pinned
        keywords. The resulting set is finally written to the DAV property.

        """

        log.info('Updating tags for %s', self.context.uniqueId)
        tms = zope.component.getUtility(zeit.retresco.interfaces.ITMS)

        if update_with is not None:
            keywords = tms.generate_keyword_list(update_with)
        else:
            keywords = tms.extract_keywords(self.context)

        E = lxml.objectify.ElementMaker(namespace=NAMESPACE)
        root = E.rankedTags()
        new_tags = E.rankedTags()
        root.append(new_tags)

        # XXX the handling of namespaces here seems chaotic
        E = lxml.objectify.ElementMaker()
        new_codes = set()

        for tag in keywords:
            if tag.code in self.disabled:
                continue
            new_codes.add(tag.code)
            new_tags.append(E.tag(tag.label, type=tag.entity_type or ''))

        old_tags = self.to_xml()
        mapped_pinned_codes = self._map_pinned_codes()
        new_pinned_codes = []
        for old_code, code in mapped_pinned_codes.iteritems():
            if code not in new_codes:
                try:
                    tag = self._find_tag_node(old_code, old_tags)
                    mapping = zeit.retresco.convert.CommonMetadata.entity_types
                    tag_type = mapping.get(tag.get('type', ''), '')
                    new_tags.append(E.tag(tag.text, type=tag_type))
                except KeyError:
                    pass
            new_pinned_codes.append(code)

        dav = zeit.connector.interfaces.IWebDAVProperties(self)
        dav[KEYWORD_PROPERTY] = lxml.etree.tostring(root.getroottree())
        dav[PINNED_PROPERTY] = "\t".join(new_pinned_codes)
        dav[DISABLED_PROPERTY] = u''
<|endoftext|>


<commit_before>from zeit.cms.repository.interfaces import ICollection, INonRecursiveCollection
import argparse
import gocept.runner
import grokcore.component as grok
import logging
import time
import transaction
import zeit.cms.celery
import zeit.cms.checkout.interfaces
import zeit.cms.content.interfaces
import zeit.cms.interfaces
import zeit.cms.repository.interfaces
import zeit.cms.workingcopy.interfaces
import zeit.retresco.interfaces
import zope.component
import zope.lifecycleevent

log = logging.getLogger(__name__)

@grok.subscribe(zope.lifecycleevent.IObjectAddedEvent)
def index_after_add(event):
    # We don't use the "extended" (object, event) method, as we are not
    # interested in the events which are dispatched to sublocations.
    context = event.object
    if not zeit.cms.interfaces.ICMSContent.providedBy(context):
        return
    if zeit.cms.repository.interfaces.IRepository.providedBy(context):
        return
    if zeit.cms.workingcopy.interfaces.IWorkingcopy.providedBy(
            event.newParent):
        return
    log.info('AfterAdd: Creating async index job for %s' % context.uniqueId)
    index_async.delay(context.uniqueId)

@grok.subscribe(
    zeit.cms.interfaces.ICMSContent,
    zeit.cms.checkout.interfaces.IAfterCheckinEvent)
def index_after_checkin(context, event):
    if event.publishing:
        # Unfortunately we have to enrich here too, even though strictly
        # speaking that "already happened" on checkin, to support the "checkin
        # and publish immediately" use case -- since there publish likely
        # happens *before* the index_async job created by checkin ran.
        index(context, enrich=True)
    else:
        index_async.apply_async((context.uniqueId,), countdown=5)

@grok.subscribe(
    zeit.cms.interfaces.ICMSContent,
    zope.lifecycleevent.IObjectRemovedEvent)
def unindex_on_remove(context, event):
    if zeit.cms.workingcopy.interfaces.IWorkingcopy.providedBy(
            event.oldParent):
        return
    unindex_async.delay(zeit.cms.content.interfaces.IUUID(context).id)

@zeit.cms.celery.task(bind=True, queuename='search')
def index_async(self, uniqueId):
    context = zeit.cms.interfaces.ICMSContent(uniqueId, None)
    if context is None:
        log.warning('Could not index %s because it does not exist any longer.',
                    uniqueId)
        return
    meta = zeit.cms.content.interfaces.ICommonMetadata(context, None)
    has_keywords = True
    if meta is not None:
        has_keywords = meta is not None and meta.keywords
    try:
        index(
            context, enrich=True, update_keywords=not has_keywords)
    except zeit.retresco.interfaces.TechnicalError:
        self.retry()

def index(content, enrich=False, update_keywords=False, publish=False):
    if update_keywords and not enrich:
        raise ValueError('enrich is required for update_keywords')
    conn = zope.component.getUtility(zeit.retresco.interfaces.ITMS)
    stack = [content]
    errors = []
    while stack:
        content = stack.pop(0)
        if (ICollection.providedBy(content) and
                not INonRecursiveCollection.providedBy(content)):
            stack.extend(content.values())
        uuid = getattr(zeit.cms.content.interfaces.IUUID(content, None), 'id',
                       '<no-uuid>')
        log.info('Updating: %s %s, enrich: %s, keywords: %s, publish: %s',
                 content.uniqueId, uuid, enrich, update_keywords, publish)
        try:
            if enrich:
                log.debug('Enriching: %s', content.uniqueId)
                response = conn.enrich(content)
                body = response.get('body')
                if update_keywords:
                    tagger = zeit.retresco.tagger.Tagger(content)
                    tagger.update(conn.generate_keyword_list(response),
                                  clear_disabled=False)
            else:
                # For reindex-only, preserve the previously enriched body.
                body = conn.get_article_data(content).get('body')

            conn.index(content, body)

            if publish:
                pub_info = zeit.cms.workflow.interfaces.IPublishInfo(content)
                if pub_info.published:
                    if zeit.retresco.interfaces.ITMSRepresentation(
                            content)() is not None:
                        log.info('Publishing: %s', content.uniqueId)
                        conn.publish(content)
                    else:
                        log.info(
                            'Skip publish for %s, missing required fields',
                            content.uniqueId)
        except zeit.retresco.interfaces.TechnicalError, e:
            log.info('Retrying %s due to %r', content.uniqueId, e)
            raise
        except Exception, e:
            errors.append(e)
            log.warning('Error indexing %s, giving up',
                        content.uniqueId, exc_info=True)
            continue
    return errors

@zeit.cms.celery.task(bind=True, queuename='search')
def unindex_async(self, uuid):
    conn = zope.component.getUtility(zeit.retresco.interfaces.ITMS)
    try:
        conn.delete_id(uuid)
    except zeit.retresco.interfaces.TechnicalError:
        self.retry()

SKIP_TYPES = ['quiz']

def should_skip(content):
    content_type = zeit.cms.type.get_type(content)
    return content_type in SKIP_TYPES

@zeit.cms.celery.task(bind=True, queuename='manual')
def index_parallel(self, unique_id, enrich=False, publish=False):
    try:
        content = zeit.cms.interfaces.ICMSContent(unique_id)
    except TypeError:
        log.warning('Could not resolve %s, giving up', unique_id)
        return
    except Exception:
        self.retry()
    if (ICollection.providedBy(content) and
            not INonRecursiveCollection.providedBy(content)):
        children = content.values()
        for item in children:
            if should_skip(item):
                log.debug('Skipping %s due to its content type', item)
                continue
            index_parallel.delay(item.uniqueId, enrich=enrich, publish=publish)
    else:
        if should_skip(content):
            log.debug('Skipping %s due to its content type', content)
            return
        start = time.time()
        try:
            errors = index(content, enrich=enrich, update_keywords=enrich,
                           publish=publish)
        except zeit.retresco.interfaces.TechnicalError:
            self.retry()
        else:
            stop = time.time()
            if not errors:
                log.info('Processed %s in %s', content.uniqueId, stop - start)

@gocept.runner.once(principal=gocept.runner.from_config(
    'zeit.retresco', 'index-principal'))
def reindex():
    parser = argparse.ArgumentParser(description='Reindex folder in TMS')
    parser.add_argument(
        'ids', type=unicode, nargs='+', help='uniqueIds to reindex')
    parser.add_argument(
        '--file', action='store_true',
        help='Load uniqueIds from a file to reindex')
    parser.add_argument(
        '--parallel', action='store_true',
        help='process via job queue instead of directly')
    parser.add_argument(
        '--enrich', action='store_true',
        help='Perform TMS analyze/enrich prior to indexing')
    parser.add_argument(
        '--publish', action='store_true',
        help='Perform TMS publish after indexing')

    args = parser.parse_args()
    ids = args.ids
    if args.file:
        if len(args.ids) > 1:
            raise Exception("Only one file can be passed!")
        with open(args.ids[0], 'r') as f:
            ids = f.read().splitlines()

    for i, id in enumerate(ids):
        if args.parallel:
            index_parallel.delay(id, args.enrich, args.publish)
            if i % 10000 == 0:
                transaction.commit()
        else:
            index(
                zeit.cms.interfaces.ICMSContent(id),
                enrich=args.enrich, update_keywords=args.enrich,
                publish=args.publish)
<commit_msg>MAINT: Document oddity<commit_after>from zeit.cms.repository.interfaces import ICollection, INonRecursiveCollection
import argparse
import gocept.runner
import grokcore.component as grok
import logging
import time
import transaction
import zeit.cms.celery
import zeit.cms.checkout.interfaces
import zeit.cms.content.interfaces
import zeit.cms.interfaces
import zeit.cms.repository.interfaces
import zeit.cms.workingcopy.interfaces
import zeit.retresco.interfaces
import zope.component
import zope.lifecycleevent

log = logging.getLogger(__name__)

@grok.subscribe(zope.lifecycleevent.IObjectAddedEvent)
def index_after_add(event):
    # We don't use the "extended" (object, event) method, as we are not
    # interested in the events which are dispatched to sublocations.
    context = event.object
    if not zeit.cms.interfaces.ICMSContent.providedBy(context):
        return
    if zeit.cms.repository.interfaces.IRepository.providedBy(context):
        return
    if zeit.cms.workingcopy.interfaces.IWorkingcopy.providedBy(
            event.newParent):
        return
    log.info('AfterAdd: Creating async index job for %s' % context.uniqueId)
    index_async.delay(context.uniqueId)

@grok.subscribe(
    zeit.cms.interfaces.ICMSContent,
    zeit.cms.checkout.interfaces.IAfterCheckinEvent)
def index_after_checkin(context, event):
    if event.publishing:
        # Unfortunately we have to enrich here too, even though strictly
        # speaking that "already happened" on checkin, to support the "checkin
        # and publish immediately" use case -- since there publish likely
        # happens *before* the index_async job created by checkin ran.
        index(context, enrich=True)
    else:
        # XXX Work around race condition between celery/redis (applies already
        # in tpc_vote) and DAV-cache in ZODB (applies only in tpc_finish, so
        # the celery job *may* start executing before that happens), BUG-796.
        index_async.apply_async((context.uniqueId,), countdown=5)

@grok.subscribe(
    zeit.cms.interfaces.ICMSContent,
    zope.lifecycleevent.IObjectRemovedEvent)
def unindex_on_remove(context, event):
    if zeit.cms.workingcopy.interfaces.IWorkingcopy.providedBy(
            event.oldParent):
        return
    unindex_async.delay(zeit.cms.content.interfaces.IUUID(context).id)

@zeit.cms.celery.task(bind=True, queuename='search')
def index_async(self, uniqueId):
    context = zeit.cms.interfaces.ICMSContent(uniqueId, None)
    if context is None:
        log.warning('Could not index %s because it does not exist any longer.',
                    uniqueId)
        return
    meta = zeit.cms.content.interfaces.ICommonMetadata(context, None)
    has_keywords = True
    if meta is not None:
        has_keywords = meta is not None and meta.keywords
    try:
        index(
            context, enrich=True, update_keywords=not has_keywords)
    except zeit.retresco.interfaces.TechnicalError:
        self.retry()

def index(content, enrich=False, update_keywords=False, publish=False):
    if update_keywords and not enrich:
        raise ValueError('enrich is required for update_keywords')
    conn = zope.component.getUtility(zeit.retresco.interfaces.ITMS)
    stack = [content]
    errors = []
    while stack:
        content = stack.pop(0)
        if (ICollection.providedBy(content) and
                not INonRecursiveCollection.providedBy(content)):
            stack.extend(content.values())
        uuid = getattr(zeit.cms.content.interfaces.IUUID(content, None), 'id',
                       '<no-uuid>')
        log.info('Updating: %s %s, enrich: %s, keywords: %s, publish: %s',
                 content.uniqueId, uuid, enrich, update_keywords, publish)
        try:
            if enrich:
                log.debug('Enriching: %s', content.uniqueId)
                response = conn.enrich(content)
                body = response.get('body')
                if update_keywords:
                    tagger = zeit.retresco.tagger.Tagger(content)
                    tagger.update(conn.generate_keyword_list(response),
                                  clear_disabled=False)
            else:
                # For reindex-only, preserve the previously enriched body.
                body = conn.get_article_data(content).get('body')

            conn.index(content, body)

            if publish:
                pub_info = zeit.cms.workflow.interfaces.IPublishInfo(content)
                if pub_info.published:
                    if zeit.retresco.interfaces.ITMSRepresentation(
                            content)() is not None:
                        log.info('Publishing: %s', content.uniqueId)
                        conn.publish(content)
                    else:
                        log.info(
                            'Skip publish for %s, missing required fields',
                            content.uniqueId)
        except zeit.retresco.interfaces.TechnicalError, e:
            log.info('Retrying %s due to %r', content.uniqueId, e)
            raise
        except Exception, e:
            errors.append(e)
            log.warning('Error indexing %s, giving up',
                        content.uniqueId, exc_info=True)
            continue
    return errors

@zeit.cms.celery.task(bind=True, queuename='search')
def unindex_async(self, uuid):
    conn = zope.component.getUtility(zeit.retresco.interfaces.ITMS)
    try:
        conn.delete_id(uuid)
    except zeit.retresco.interfaces.TechnicalError:
        self.retry()

SKIP_TYPES = ['quiz']

def should_skip(content):
    content_type = zeit.cms.type.get_type(content)
    return content_type in SKIP_TYPES

@zeit.cms.celery.task(bind=True, queuename='manual')
def index_parallel(self, unique_id, enrich=False, publish=False):
    try:
        content = zeit.cms.interfaces.ICMSContent(unique_id)
    except TypeError:
        log.warning('Could not resolve %s, giving up', unique_id)
        return
    except Exception:
        self.retry()
    if (ICollection.providedBy(content) and
            not INonRecursiveCollection.providedBy(content)):
        children = content.values()
        for item in children:
            if should_skip(item):
                log.debug('Skipping %s due to its content type', item)
                continue
            index_parallel.delay(item.uniqueId, enrich=enrich, publish=publish)
    else:
        if should_skip(content):
            log.debug('Skipping %s due to its content type', content)
            return
        start = time.time()
        try:
            errors = index(content, enrich=enrich, update_keywords=enrich,
                           publish=publish)
        except zeit.retresco.interfaces.TechnicalError:
            self.retry()
        else:
            stop = time.time()
            if not errors:
                log.info('Processed %s in %s', content.uniqueId, stop - start)

@gocept.runner.once(principal=gocept.runner.from_config(
    'zeit.retresco', 'index-principal'))
def reindex():
    parser = argparse.ArgumentParser(description='Reindex folder in TMS')
    parser.add_argument(
        'ids', type=unicode, nargs='+', help='uniqueIds to reindex')
    parser.add_argument(
        '--file', action='store_true',
        help='Load uniqueIds from a file to reindex')
    parser.add_argument(
        '--parallel', action='store_true',
        help='process via job queue instead of directly')
    parser.add_argument(
        '--enrich', action='store_true',
        help='Perform TMS analyze/enrich prior to indexing')
    parser.add_argument(
        '--publish', action='store_true',
        help='Perform TMS publish after indexing')

    args = parser.parse_args()
    ids = args.ids
    if args.file:
        if len(args.ids) > 1:
            raise Exception("Only one file can be passed!")
        with open(args.ids[0], 'r') as f:
            ids = f.read().splitlines()

    for i, id in enumerate(ids):
        if args.parallel:
            index_parallel.delay(id, args.enrich, args.publish)
            if i % 10000 == 0:
                transaction.commit()
        else:
            index(
                zeit.cms.interfaces.ICMSContent(id),
                enrich=args.enrich, update_keywords=args.enrich,
                publish=args.publish)
<|endoftext|>


<commit_before>import os.path as op
import pandas as pd
import numpy as np
from collections import defaultdict
from Bio import AlignIO
from Bio.Emboss.Applications import NeedleCommandline

def run_needle_alignment_on_files(id_a, faa_a, id_b, faa_b, gapopen=10, gapextend=0.5, outdir='', outfile=''):
    """Run the needle alignment program for two fasta files and writes the result to a file. Returns the filename.

    Args:
        id_a: first id
        faa_a: first FASTA file path
        id_b: second id
        faa_b: second FASTA file path

    Returns:
        str: alignment_file - file path to alignment

    """

    if not outfile:
        outfile = op.join(outdir, '{}_{}_align.txt'.format(id_a, id_b))
    else:
        outfile = op.join(outdir, outfile)

    if op.exists(outfile):
        return outfile

    needle_cline = NeedleCommandline(asequence=faa_a, bsequence=faa_b,
                                     gapopen=gapopen, gapextend=gapextend,
                                     outfile=outfile)
    stdout, stderr = needle_cline()
    return outfile

def run_needle_alignment_on_str(id_a, seq_a, id_b, seq_b, gapopen=10, gapextend=0.5, outdir='', outfile=''):
    """Run the needle alignment program and writes the result to a file. Returns the filename.

    SEE: https://www.biostars.org/p/91124/

    Args:
        id_a (str): sequence ID #1
        seq_a (str): sequence #1
        id_b (str): sequence ID #2
        seq_b (str): sequence #2

    Returns:
        alignment_file - file name of alignment

    """

    if not outfile:
        outfile = op.join(outdir, '{}_{}_align.txt'.format(id_a, id_b))
    else:
        outfile = op.join(outdir, outfile)

    if op.exists(outfile):
        return outfile

    needle_cline = NeedleCommandline(asequence="asis::"+seq_a, bsequence="asis::"+seq_b,
                                     gapopen=gapopen, gapextend=gapextend,
                                     outfile=outfile)
    stdout, stderr = needle_cline()

    return outfile

def get_alignment_summary_df(alignment_file):
    """Create an alignment dataframe indicating where matches, mutations, insertions, and deletions are.
        Takes in a needle alignment file and returns a Pandas DataFrame.

    Args:
        alignment_file - needle alignment file between two amino acid sequences

    Retunrs:
        alignment_df - Pandas DataFrame of the alignment.

    """

    alignments = list(AlignIO.parse(alignment_file, "emboss"))
    alignment_df = pd.DataFrame(columns=['id_a', 'id_b', 'type', 'id_a_start', 'id_a_stop', 'count', 'id_a_aa', 'id_b_aa'])

    for alignment in alignments:
        a_seq_id = list(alignment)[0].id
        a_seq = str(list(alignment)[0].seq)
        b_seq_id = list(alignment)[1].id
        b_seq = str(list(alignment)[1].seq)

        idx_start = 1

        uniprot_shift = 0
        insertion_counter = 0
        for i, (a, b) in enumerate(zip(a_seq, b_seq)):
            if uniprot_shift != 0:
                new_i = i - uniprot_shift
            else:
                new_i = i

            if a == b and a != '-' and b != '-':
                aa_flag = 'match'
            if a != b and a == '-' and b != '-':
                aa_flag = 'insertion'
            if a != b and a != '-' and b == '-':
                aa_flag = 'deletion'
            if a != b and a != '-' and b == 'X':
                aa_flag = 'unresolved'
            elif a != b and a != '-' and b != '-':
                aa_flag = 'mutation'
            if i == 0:
                aa_flag_tracker = aa_flag

            if aa_flag == 'insertion':
                uniprot_shift += 1
                insertion_counter += 1

            if aa_flag != aa_flag_tracker:
                idx_end = new_i
                if aa_flag_tracker == 'match' or aa_flag_tracker == 'deletion' or aa_flag_tracker == 'unresolved':
                    appender = {}
                    appender['id_a'] = a_seq_id
                    appender['id_b'] = b_seq_id
                    appender['type'] = aa_flag_tracker
                    appender['id_a_start'] = idx_start
                    appender['id_a_stop'] = idx_end
                    appender['count'] = idx_end - idx_start + 1
                    alignment_df = alignment_df.append(appender, ignore_index=True)

                if aa_flag_tracker == 'insertion':
                    appender = {}
                    appender['id_a'] = a_seq_id
                    appender['id_b'] = b_seq_id
                    appender['type'] = aa_flag_tracker
                    appender['id_a_start'] = new_i
                    appender['id_a_stop'] = new_i
                    appender['count'] = insertion_counter
                    alignment_df = alignment_df.append(appender, ignore_index=True)
                    insertion_counter = 0
                idx_start = new_i + 1

            # TODO: test if this should be outside the for loop
            if aa_flag == 'mutation':
                appender = {}
                appender['id_a'] = a_seq_id
                appender['id_b'] = b_seq_id
                appender['type'] = aa_flag
                appender['id_a_start'] = new_i + 1
                appender['id_a_stop'] = new_i + 1
                appender['count'] = 1
                appender['id_a_aa'] = a
                appender['id_b_aa'] = b
                alignment_df = alignment_df.append(appender, ignore_index=True)
                idx_start = new_i + 1

            elif (i + 1) == len(a_seq) and aa_flag_tracker == aa_flag:
                idx_end = new_i + 1
                if aa_flag_tracker != 'mutation' and aa_flag_tracker != 'insertion':
                    appender = {}
                    appender['id_a'] = a_seq_id
                    appender['id_b'] = b_seq_id
                    appender['type'] = aa_flag_tracker
                    appender['id_a_start'] = idx_start
                    appender['id_a_stop'] = idx_end
                    appender['count'] = idx_end - idx_start + 1
                    alignment_df = alignment_df.append(appender, ignore_index=True)

            elif (i + 1) == len(a_seq) and aa_flag_tracker != aa_flag:
                idx_end = new_i + 1
                idx_start = new_i + 1
                if aa_flag_tracker != 'mutation' and aa_flag_tracker != 'insertion':
                    appender = {}
                    appender['id_a'] = a_seq_id
                    appender['id_b'] = b_seq_id
                    appender['type'] = aa_flag
                    appender['id_a_start'] = idx_start
                    appender['id_a_stop'] = idx_end
                    appender['count'] = idx_end - idx_start + 1
                    alignment_df = alignment_df.append(appender, ignore_index=True)

            aa_flag_tracker = aa_flag

    return alignment_df

def get_alignment_df(alignment_file, a_seq_id=None, b_seq_id=None):
    """Get a Pandas DataFrame of the Needle alignment results. Contains all positions of the sequences.

    Args:
        alignment_file:
        a_seq_id: Optional specification of the ID of the reference sequence
        b_seq_id: Optional specification of the ID of the aligned sequence

    Returns:
        Pandas DataFrame: all positions in the alignment

    """
    alignments = list(AlignIO.parse(alignment_file, "emboss"))

    appender = defaultdict(dict)
    idx = 0
    for alignment in alignments:
        if not a_seq_id:
            a_seq_id = list(alignment)[0].id
        a_seq = str(list(alignment)[0].seq)
        if not b_seq_id:
            b_seq_id = list(alignment)[1].id
        b_seq = str(list(alignment)[1].seq)

        a_idx = 1
        b_idx = 1

        for i, (a,b) in enumerate(zip(a_seq,b_seq)):
            if a == b and a != '-' and b != '-':
                aa_flag = 'match'
            if a != b and a == '-' and b != '-':
                aa_flag = 'insertion'
            if a != b and a != '-' and b == '-':
                aa_flag = 'deletion'
            if a != b and a != '-' and b == 'X':
                aa_flag = 'unresolved'
            if a != b and b != '-' and a == 'X':
                aa_flag = 'unresolved'
            elif a != b and a != '-' and b != '-':
                aa_flag = 'mutation'

            appender[idx]['id_a'] = a_seq_id
            appender[idx]['id_b'] = b_seq_id
            appender[idx]['type'] = aa_flag

            if aa_flag == 'match' or aa_flag == 'unresolved' or aa_flag == 'mutation':
                appender[idx]['id_a_aa'] = a
                appender[idx]['id_a_pos'] = a_idx
                appender[idx]['id_b_aa'] = b
                appender[idx]['id_b_pos'] = b_idx
                a_idx += 1
                b_idx += 1

            if aa_flag == 'deletion':
                appender[idx]['id_a_aa'] = a
                appender[idx]['id_a_pos'] = a_idx
                a_idx += 1

            if aa_flag == 'insertion':
                appender[idx]['id_b_aa'] = b
                appender[idx]['id_b_pos'] = b_idx
                b_idx += 1

            idx += 1

    alignment_df = pd.DataFrame.from_dict(appender, orient='index')
    alignment_df = alignment_df[['id_a', 'id_b', 'type', 'id_a_aa', 'id_a_pos', 'id_b_aa', 'id_b_pos']].fillna(value=np.nan)

    return alignment_df

def needle_statistics(infile):
    """Reads in a needle alignment file and spits out statistics of the alignment.

    Args:
        infile (str): Alignment file name

    Returns:
        dict: alignment_properties - a dictionary telling you the number of gaps, identity, etc.

    """

    alignments = list(AlignIO.parse(infile, "emboss"))

    f = open(infile)
    alignment_properties = defaultdict(dict)
    line = f.readline()

    for i in range(len(alignments)):
        while line.rstrip() != "#=======================================":
            line = f.readline()
            if not line:
                raise StopIteration

        while line[0] == "#":
            # Read in the rest of this alignment header,
            # try and discover the number of records expected and their length
            parts = line[1:].split(":", 1)
            key = parts[0].lower().strip()
            if key == '2':
                pdb_id = parts[1].strip()
            if key == 'identity':
                ident_parse = parts[1].strip().replace('(','').replace(')','').replace('%','').split()
                ident_num = int(ident_parse[0].split('/')[0])
                ident_percent = float(ident_parse[1])
                alignment_properties[pdb_id]['identity'] = ident_num
                alignment_properties[pdb_id]['identity_percent'] = ident_percent
            if key == 'similarity':
                sim_parse = parts[1].strip().replace('(','').replace(')','').replace('%','').split()
                sim_num = int(sim_parse[0].split('/')[0])
                sim_percent = float(sim_parse[1])
                alignment_properties[pdb_id]['similarity'] = sim_num
                alignment_properties[pdb_id]['similarity_percent'] = sim_percent
            if key == 'gaps':
                gap_parse = parts[1].strip().replace('(','').replace(')','').replace('%','').split()
                gap_num = int(gap_parse[0].split('/')[0])
                gap_percent = float(gap_parse[1])
                alignment_properties[pdb_id]['gaps'] = gap_num
                alignment_properties[pdb_id]['gaps_percent'] = gap_percent
            if key == 'score':
                score = float(parts[1].strip())
                alignment_properties[pdb_id]['score'] = score

            # And read in another line...
            line = f.readline()

    return alignment_properties

# TODO: needleall has a bug
# def run_alignment_needleall(a_id, a_faa, b_id, b_faa):
#
#     from Bio.Emboss.Applications import NeedleallCommandline
#     import os.path
#
#     alignment_file = "%s-%s_align.txt" % (a_id, b_id)
#
#     if os.path.isfile(alignment_file):
#         # print 'Alignment %s file already exists' % alignment_file
#         return alignment_file
#
#     else:
#         print '**RUNNING ALIGNMENT FOR %kegg AND %kegg**' % (a_id, b_id)
#         needle_cline = NeedleallCommandline(asequence=a_faa, bsequence=b_faa, gapopen=10, gapextend=0.5, outfile=alignment_file)
#         stdout, stderr = needle_cline()
#         return alignment_file
<commit_msg>Needle alignment now returns raw results in string form, also loads cached results<commit_after>import os.path as op
import pandas as pd
import numpy as np
from collections import defaultdict
from Bio import AlignIO
from Bio.Emboss.Applications import NeedleCommandline

def run_needle_alignment_on_files(id_a, faa_a, id_b, faa_b, gapopen=10, gapextend=0.5,
                                  write_output=False, outdir='', outfile='', force_rerun=False):
    """Run the needle alignment program for two fasta files and return the raw alignment result.

    More info:
    EMBOSS needle: http://www.bioinformatics.nl/cgi-bin/emboss/help/needle
    Biopython wrapper: http://biopython.org/DIST/docs/tutorial/Tutorial.html#htoc84

    Args:
        id_a: ID of reference sequence
        faa_a: File path to reference sequence
        id_b: ID of sequence to be aligned
        faa_b: File path to sequence to be aligned
        gapopen: Gap open penalty is the score taken away when a gap is created
        gapextend: Gap extension penalty is added to the standard gap penalty for each base or residue in the gap
        write_output (bool): Default False, set to True if you want the alignment file saved
        outdir (str, optional): Path to output directory. Default is the current directory.
        outfile (str, optional): Name of output file. If not set, is {id_a}_{id_b}_align.txt
        force_rerun (bool): Default False, set to True if you want to rerun the alignment if outfile exists.

    Returns:
        str: Raw alignment result of the needle alignment in srspair format.

    """

    # If you don't want to save the output file, just run the alignment and return the raw results
    if not write_output:
        needle_cline = NeedleCommandline(asequence=faa_a, bsequence=faa_b,
                                         gapopen=gapopen, gapextend=gapextend,
                                         stdout=True, auto=True)
        raw_alignment_text, stderr = needle_cline()

    # If you do want to save an output file...
    else:
        # Make a default name if no outfile is set
        if not outfile:
            outfile = op.join(outdir, '{}_{}_align.txt'.format(id_a, id_b))
        else:
            outfile = op.join(outdir, outfile)

        # Check if the outfile already exists, and read it and return those results
        if op.exists(outfile) and not force_rerun:
            with open(outfile) as f:
                raw_alignment_text = f.read()

        # If it doesn't exist, or force_rerun=True, run the alignment
        else:
            needle_cline = NeedleCommandline(asequence=faa_a, bsequence=faa_b,
                                             gapopen=gapopen, gapextend=gapextend,
                                             outfile=outfile)
            stdout, stderr = needle_cline()
            with open(outfile) as f:
                raw_alignment_text = f.read()

    return raw_alignment_text

def run_needle_alignment_on_str(id_a, seq_a, id_b, seq_b, gapopen=10, gapextend=0.5,
                                write_output=False, outdir='', outfile='', force_rerun=False):
    """Run the needle alignment program for two strings and return the raw alignment result.

    More info:
    EMBOSS needle: http://www.bioinformatics.nl/cgi-bin/emboss/help/needle
    Biopython wrapper: http://biopython.org/DIST/docs/tutorial/Tutorial.html#htoc84
    Using strings as input: https://www.biostars.org/p/91124/

    Args:
        id_a: ID of reference sequence
        seq_a: String representation of reference sequence
        id_b: ID of sequence to be aligned
        seq_b: String representation of sequence to be aligned
        gapopen: Gap open penalty is the score taken away when a gap is created
        gapextend: Gap extension penalty is added to the standard gap penalty for each base or residue in the gap
        write_output (bool): Default False, set to True if you want the alignment file saved
        outdir (str, optional): Path to output directory. Default is the current directory.
        outfile (str, optional): Name of output file. If not set, is {id_a}_{id_b}_align.txt
        force_rerun (bool): Default False, set to True if you want to rerun the alignment if outfile exists.

    Returns:
        str: Raw alignment result of the needle alignment in srspair format.

    """
    # If you don't want to save the output file, just run the alignment and return the raw results
    if not write_output:
        needle_cline = NeedleCommandline(asequence="asis::"+seq_a, bsequence="asis::"+seq_b,
                                         gapopen=gapopen, gapextend=gapextend,
                                         stdout=True, auto=True)
        raw_alignment_text, stderr = needle_cline()

    # If you do want to save an output file...
    else:
        # Make a default name if no outfile is set
        if not outfile:
            outfile = op.join(outdir, '{}_{}_align.txt'.format(id_a, id_b))
        else:
            outfile = op.join(outdir, outfile)

        # Check if the outfile already exists, and read it and return those results
        if op.exists(outfile) and not force_rerun:
            with open(outfile) as f:
                raw_alignment_text = f.read()

        # If it doesn't exist, or force_rerun=True, run the alignment
        else:
            needle_cline = NeedleCommandline(asequence="asis::"+seq_a, bsequence="asis::"+seq_b,
                                             gapopen=gapopen, gapextend=gapextend,
                                             outfile=outfile)
            stdout, stderr = needle_cline()
            with open(outfile) as f:
                raw_alignment_text = f.read()

    return raw_alignment_text

def get_alignment_summary_df(alignment_file):
    """Create an alignment dataframe indicating where matches, mutations, insertions, and deletions are.
        Takes in a needle alignment file and returns a Pandas DataFrame.

    Args:
        alignment_file - needle alignment file between two amino acid sequences

    Retunrs:
        alignment_df - Pandas DataFrame of the alignment.

    """

    alignments = list(AlignIO.parse(alignment_file, "emboss"))
    alignment_df = pd.DataFrame(columns=['id_a', 'id_b', 'type', 'id_a_start', 'id_a_stop', 'count', 'id_a_aa', 'id_b_aa'])

    for alignment in alignments:
        a_seq_id = list(alignment)[0].id
        a_seq = str(list(alignment)[0].seq)
        b_seq_id = list(alignment)[1].id
        b_seq = str(list(alignment)[1].seq)

        idx_start = 1

        uniprot_shift = 0
        insertion_counter = 0
        for i, (a, b) in enumerate(zip(a_seq, b_seq)):
            if uniprot_shift != 0:
                new_i = i - uniprot_shift
            else:
                new_i = i

            if a == b and a != '-' and b != '-':
                aa_flag = 'match'
            if a != b and a == '-' and b != '-':
                aa_flag = 'insertion'
            if a != b and a != '-' and b == '-':
                aa_flag = 'deletion'
            if a != b and a != '-' and b == 'X':
                aa_flag = 'unresolved'
            elif a != b and a != '-' and b != '-':
                aa_flag = 'mutation'
            if i == 0:
                aa_flag_tracker = aa_flag

            if aa_flag == 'insertion':
                uniprot_shift += 1
                insertion_counter += 1

            if aa_flag != aa_flag_tracker:
                idx_end = new_i
                if aa_flag_tracker == 'match' or aa_flag_tracker == 'deletion' or aa_flag_tracker == 'unresolved':
                    appender = {}
                    appender['id_a'] = a_seq_id
                    appender['id_b'] = b_seq_id
                    appender['type'] = aa_flag_tracker
                    appender['id_a_start'] = idx_start
                    appender['id_a_stop'] = idx_end
                    appender['count'] = idx_end - idx_start + 1
                    alignment_df = alignment_df.append(appender, ignore_index=True)

                if aa_flag_tracker == 'insertion':
                    appender = {}
                    appender['id_a'] = a_seq_id
                    appender['id_b'] = b_seq_id
                    appender['type'] = aa_flag_tracker
                    appender['id_a_start'] = new_i
                    appender['id_a_stop'] = new_i
                    appender['count'] = insertion_counter
                    alignment_df = alignment_df.append(appender, ignore_index=True)
                    insertion_counter = 0
                idx_start = new_i + 1

            # TODO: test if this should be outside the for loop
            if aa_flag == 'mutation':
                appender = {}
                appender['id_a'] = a_seq_id
                appender['id_b'] = b_seq_id
                appender['type'] = aa_flag
                appender['id_a_start'] = new_i + 1
                appender['id_a_stop'] = new_i + 1
                appender['count'] = 1
                appender['id_a_aa'] = a
                appender['id_b_aa'] = b
                alignment_df = alignment_df.append(appender, ignore_index=True)
                idx_start = new_i + 1

            elif (i + 1) == len(a_seq) and aa_flag_tracker == aa_flag:
                idx_end = new_i + 1
                if aa_flag_tracker != 'mutation' and aa_flag_tracker != 'insertion':
                    appender = {}
                    appender['id_a'] = a_seq_id
                    appender['id_b'] = b_seq_id
                    appender['type'] = aa_flag_tracker
                    appender['id_a_start'] = idx_start
                    appender['id_a_stop'] = idx_end
                    appender['count'] = idx_end - idx_start + 1
                    alignment_df = alignment_df.append(appender, ignore_index=True)

            elif (i + 1) == len(a_seq) and aa_flag_tracker != aa_flag:
                idx_end = new_i + 1
                idx_start = new_i + 1
                if aa_flag_tracker != 'mutation' and aa_flag_tracker != 'insertion':
                    appender = {}
                    appender['id_a'] = a_seq_id
                    appender['id_b'] = b_seq_id
                    appender['type'] = aa_flag
                    appender['id_a_start'] = idx_start
                    appender['id_a_stop'] = idx_end
                    appender['count'] = idx_end - idx_start + 1
                    alignment_df = alignment_df.append(appender, ignore_index=True)

            aa_flag_tracker = aa_flag

    return alignment_df

def get_alignment_df(alignment_file, a_seq_id=None, b_seq_id=None):
    """Get a Pandas DataFrame of the Needle alignment results. Contains all positions of the sequences.

    Args:
        alignment_file:
        a_seq_id: Optional specification of the ID of the reference sequence
        b_seq_id: Optional specification of the ID of the aligned sequence

    Returns:
        Pandas DataFrame: all positions in the alignment

    """
    alignments = list(AlignIO.parse(alignment_file, "emboss"))

    appender = defaultdict(dict)
    idx = 0
    for alignment in alignments:
        if not a_seq_id:
            a_seq_id = list(alignment)[0].id
        a_seq = str(list(alignment)[0].seq)
        if not b_seq_id:
            b_seq_id = list(alignment)[1].id
        b_seq = str(list(alignment)[1].seq)

        a_idx = 1
        b_idx = 1

        for i, (a,b) in enumerate(zip(a_seq,b_seq)):
            if a == b and a != '-' and b != '-':
                aa_flag = 'match'
            if a != b and a == '-' and b != '-':
                aa_flag = 'insertion'
            if a != b and a != '-' and b == '-':
                aa_flag = 'deletion'
            if a != b and a != '-' and b == 'X':
                aa_flag = 'unresolved'
            if a != b and b != '-' and a == 'X':
                aa_flag = 'unresolved'
            elif a != b and a != '-' and b != '-':
                aa_flag = 'mutation'

            appender[idx]['id_a'] = a_seq_id
            appender[idx]['id_b'] = b_seq_id
            appender[idx]['type'] = aa_flag

            if aa_flag == 'match' or aa_flag == 'unresolved' or aa_flag == 'mutation':
                appender[idx]['id_a_aa'] = a
                appender[idx]['id_a_pos'] = a_idx
                appender[idx]['id_b_aa'] = b
                appender[idx]['id_b_pos'] = b_idx
                a_idx += 1
                b_idx += 1

            if aa_flag == 'deletion':
                appender[idx]['id_a_aa'] = a
                appender[idx]['id_a_pos'] = a_idx
                a_idx += 1

            if aa_flag == 'insertion':
                appender[idx]['id_b_aa'] = b
                appender[idx]['id_b_pos'] = b_idx
                b_idx += 1

            idx += 1

    alignment_df = pd.DataFrame.from_dict(appender, orient='index')
    alignment_df = alignment_df[['id_a', 'id_b', 'type', 'id_a_aa', 'id_a_pos', 'id_b_aa', 'id_b_pos']].fillna(value=np.nan)

    return alignment_df

def needle_statistics(infile):
    """Reads in a needle alignment file and spits out statistics of the alignment.

    Args:
        infile (str): Alignment file name

    Returns:
        dict: alignment_properties - a dictionary telling you the number of gaps, identity, etc.

    """

    alignments = list(AlignIO.parse(infile, "emboss"))

    f = open(infile)
    alignment_properties = defaultdict(dict)
    line = f.readline()

    for i in range(len(alignments)):
        while line.rstrip() != "#=======================================":
            line = f.readline()
            if not line:
                raise StopIteration

        while line[0] == "#":
            # Read in the rest of this alignment header,
            # try and discover the number of records expected and their length
            parts = line[1:].split(":", 1)
            key = parts[0].lower().strip()
            if key == '2':
                pdb_id = parts[1].strip()
            if key == 'identity':
                ident_parse = parts[1].strip().replace('(','').replace(')','').replace('%','').split()
                ident_num = int(ident_parse[0].split('/')[0])
                ident_percent = float(ident_parse[1])
                alignment_properties[pdb_id]['identity'] = ident_num
                alignment_properties[pdb_id]['identity_percent'] = ident_percent
            if key == 'similarity':
                sim_parse = parts[1].strip().replace('(','').replace(')','').replace('%','').split()
                sim_num = int(sim_parse[0].split('/')[0])
                sim_percent = float(sim_parse[1])
                alignment_properties[pdb_id]['similarity'] = sim_num
                alignment_properties[pdb_id]['similarity_percent'] = sim_percent
            if key == 'gaps':
                gap_parse = parts[1].strip().replace('(','').replace(')','').replace('%','').split()
                gap_num = int(gap_parse[0].split('/')[0])
                gap_percent = float(gap_parse[1])
                alignment_properties[pdb_id]['gaps'] = gap_num
                alignment_properties[pdb_id]['gaps_percent'] = gap_percent
            if key == 'score':
                score = float(parts[1].strip())
                alignment_properties[pdb_id]['score'] = score

            # And read in another line...
            line = f.readline()

    return alignment_properties

# TODO: needleall has a bug
# def run_alignment_needleall(a_id, a_faa, b_id, b_faa):
#
#     from Bio.Emboss.Applications import NeedleallCommandline
#     import os.path
#
#     alignment_file = "%s-%s_align.txt" % (a_id, b_id)
#
#     if os.path.isfile(alignment_file):
#         # print 'Alignment %s file already exists' % alignment_file
#         return alignment_file
#
#     else:
#         print '**RUNNING ALIGNMENT FOR %kegg AND %kegg**' % (a_id, b_id)
#         needle_cline = NeedleallCommandline(asequence=a_faa, bsequence=b_faa, gapopen=10, gapextend=0.5, outfile=alignment_file)
#         stdout, stderr = needle_cline()
#         return alignment_file
<|endoftext|>


<commit_before>#!/usr/bin/env python
"""
SCL; 26 Jul 2012.
"""

from StringIO import StringIO
import copy
from tulip.automaton import Automaton, AutomatonState

REFERENCE_AUTFILE = """
smv file: /Users/scott/scott_centre/scm/tulip-xml-branch/examples/specs/robot_simple.smv
spc file: /Users/scott/scott_centre/scm/tulip-xml-branch/examples/specs/robot_simple.spc
priority kind: 3
Specification is realizable...
==== Building an implementation =========
-----------------------------------------
State 0 with rank 0 -> <park:1, cellID:0, X0reach:0>
	With successors : 1, 2
State 1 with rank 1 -> <park:0, cellID:0, X0reach:0>
	With successors : 3, 4
State 2 with rank 1 -> <park:1, cellID:0, X0reach:0>
	With successors : 3, 4
State 3 with rank 1 -> <park:0, cellID:1, X0reach:0>
	With successors : 5, 6
State 4 with rank 1 -> <park:1, cellID:1, X0reach:0>
	With successors : 5, 6
State 5 with rank 1 -> <park:0, cellID:4, X0reach:0>
	With successors : 7, 6
State 6 with rank 1 -> <park:1, cellID:4, X0reach:0>
	With successors : 7, 6
State 7 with rank 1 -> <park:0, cellID:3, X0reach:0>
	With successors : 8, 9
State 8 with rank 1 -> <park:0, cellID:4, X0reach:1>
	With successors : 15, 16
State 9 with rank 1 -> <park:1, cellID:4, X0reach:1>
	With successors : 10, 11
State 10 with rank 0 -> <park:0, cellID:4, X0reach:0>
	With successors : 12, 13
State 11 with rank 0 -> <park:1, cellID:4, X0reach:0>
	With successors : 12, 13
State 12 with rank 0 -> <park:0, cellID:1, X0reach:0>
	With successors : 14, 0
State 13 with rank 0 -> <park:1, cellID:1, X0reach:0>
	With successors : 14, 0
State 14 with rank 0 -> <park:0, cellID:0, X0reach:0>
	With successors : 1, 2, 1, 2
State 15 with rank 0 -> <park:0, cellID:4, X0reach:1>
	With successors : 17, 18
State 16 with rank 0 -> <park:1, cellID:4, X0reach:1>
	With successors : 12, 13
State 17 with rank 0 -> <park:0, cellID:1, X0reach:1>
	With successors : 19, 20
State 18 with rank 0 -> <park:1, cellID:1, X0reach:1>
	With successors : 14, 0
State 19 with rank 0 -> <park:0, cellID:0, X0reach:1>
	With successors : 21, 22
State 20 with rank 0 -> <park:1, cellID:0, X0reach:1>
	With successors : 1, 2
State 21 with rank 1 -> <park:0, cellID:0, X0reach:1>
	With successors : 19, 20
State 22 with rank 1 -> <park:1, cellID:0, X0reach:1>
	With successors : 14, 0
-----------------------------------------
Games time: 12
Checking realizability time: 14
Strategy time: 240
"""

def automaton2xml_test():
    ref_aut = Automaton(states_or_file=StringIO(REFERENCE_AUTFILE))
    aut = Automaton()
    aut.loadXML(ref_aut.dumpXML())
    assert aut == ref_aut
    aut.states[0].transition = []
    assert aut != ref_aut

def node_copy_test():
    node1 = AutomatonState(id=0, state={"nyan":0, "cat":1}, transition=[0])
    node2 = node1.copy()
    assert node2 is not node1
    assert node2.id == 0
    assert node2.state == {"nyan":0, "cat":1}
    assert node2.transition == [0]

class basic_Automaton_test():
    def setUp(self):
        self.empty_aut = Automaton()
        self.singleton = Automaton(states_or_file=[AutomatonState(id=0, state={"x":0}, transition=[0])])
        self.chain_len = 8
        base_state = dict([("x"+str(k),0) for k in range(self.chain_len)])
        self.chain_aut = Automaton()
        for k in range(self.chain_len):
            base_state["x"+str(k)] = 1
            self.chain_aut.addAutState(AutomatonState(id=k, state=base_state, transition=[(k+1)%self.chain_len]))
            base_state["x"+str(k)] = 0

    def tearDown(self):
        pass

    def test_equality(self):
        assert self.empty_aut == self.empty_aut
        assert self.singleton == self.singleton
        assert self.empty_aut != self.singleton

    def test_size(self):
        assert len(self.empty_aut) == 0  # Should give same result as size()
        assert self.empty_aut.size() == 0
        assert self.singleton.size() == 1
        assert len(self.chain_aut) == self.chain_len

    def test_copy(self):
        A = self.singleton.copy()
        assert len(A) == 1
        B = copy.copy(self.singleton)
        assert A == B
        assert A is not B

        C = self.chain_aut.copy()
        assert len(C) == self.chain_len
        assert C == self.chain_aut
        base_state = dict([("x"+str(k),0) for k in range(self.chain_len)])
        base_state["x0"] = 1
        assert C.findAutState(base_state) != -1
        C.findAutState(base_state).state["x0"] = 0
        assert C != self.chain_aut
<commit_msg>TEST: Automaton load/dump XML methods.  The test case is small and simple.<commit_after>#!/usr/bin/env python
"""
SCL; 15 August 2012.
"""

from StringIO import StringIO
import copy
from tulip.automaton import Automaton, AutomatonState

# Useful for regression testing; may be good to later provide
# backwards-compatibility testing for prior tulipcon versions.
REFERENCE_XMLFRAGMENT = """  <aut>
    <node>
      <id>0</id><name></name>
      <child_list>1 2</child_list>
      <state><item key="y" value="0" /><item key="x" value="1" /></state>
    </node>
    <node>
      <id>1</id><name></name>
      <child_list>1 2</child_list>
      <state><item key="y" value="0" /><item key="x" value="0" /></state>
    </node>
    <node>
      <id>2</id><name></name>
      <child_list>1 0</child_list>
      <state><item key="y" value="1" /><item key="x" value="1" /></state>
    </node>
  </aut>
"""

REFERENCE_AUTFILE = """
smv file: /Users/scott/scott_centre/scm/tulip-xml-branch/examples/specs/robot_simple.smv
spc file: /Users/scott/scott_centre/scm/tulip-xml-branch/examples/specs/robot_simple.spc
priority kind: 3
Specification is realizable...
==== Building an implementation =========
-----------------------------------------
State 0 with rank 0 -> <park:1, cellID:0, X0reach:0>
	With successors : 1, 2
State 1 with rank 1 -> <park:0, cellID:0, X0reach:0>
	With successors : 3, 4
State 2 with rank 1 -> <park:1, cellID:0, X0reach:0>
	With successors : 3, 4
State 3 with rank 1 -> <park:0, cellID:1, X0reach:0>
	With successors : 5, 6
State 4 with rank 1 -> <park:1, cellID:1, X0reach:0>
	With successors : 5, 6
State 5 with rank 1 -> <park:0, cellID:4, X0reach:0>
	With successors : 7, 6
State 6 with rank 1 -> <park:1, cellID:4, X0reach:0>
	With successors : 7, 6
State 7 with rank 1 -> <park:0, cellID:3, X0reach:0>
	With successors : 8, 9
State 8 with rank 1 -> <park:0, cellID:4, X0reach:1>
	With successors : 15, 16
State 9 with rank 1 -> <park:1, cellID:4, X0reach:1>
	With successors : 10, 11
State 10 with rank 0 -> <park:0, cellID:4, X0reach:0>
	With successors : 12, 13
State 11 with rank 0 -> <park:1, cellID:4, X0reach:0>
	With successors : 12, 13
State 12 with rank 0 -> <park:0, cellID:1, X0reach:0>
	With successors : 14, 0
State 13 with rank 0 -> <park:1, cellID:1, X0reach:0>
	With successors : 14, 0
State 14 with rank 0 -> <park:0, cellID:0, X0reach:0>
	With successors : 1, 2, 1, 2
State 15 with rank 0 -> <park:0, cellID:4, X0reach:1>
	With successors : 17, 18
State 16 with rank 0 -> <park:1, cellID:4, X0reach:1>
	With successors : 12, 13
State 17 with rank 0 -> <park:0, cellID:1, X0reach:1>
	With successors : 19, 20
State 18 with rank 0 -> <park:1, cellID:1, X0reach:1>
	With successors : 14, 0
State 19 with rank 0 -> <park:0, cellID:0, X0reach:1>
	With successors : 21, 22
State 20 with rank 0 -> <park:1, cellID:0, X0reach:1>
	With successors : 1, 2
State 21 with rank 1 -> <park:0, cellID:0, X0reach:1>
	With successors : 19, 20
State 22 with rank 1 -> <park:1, cellID:0, X0reach:1>
	With successors : 14, 0
-----------------------------------------
Games time: 12
Checking realizability time: 14
Strategy time: 240
"""

def automaton2xml_test():
    ref_aut = Automaton(states_or_file=StringIO(REFERENCE_AUTFILE))
    aut = Automaton()
    aut.loadXML(ref_aut.dumpXML())
    assert aut == ref_aut
    aut.states[0].transition = []
    assert aut != ref_aut

def loadXML_test():
    aut = Automaton()
    aut.loadXML(REFERENCE_XMLFRAGMENT)
    assert len(aut) == 3
    assert (len(aut.states[0].state.keys()) == 2) and ("x" in aut.states[0].state.keys()) and ("y" in aut.states[0].state.keys())
    node0 = aut.findAutState({"x": 1, "y": 0})
    node1 = aut.findAutState({"x": 0, "y": 0})
    node2 = aut.findAutState({"x": 1, "y": 1})
    assert (len(node0.transition) == 2) and (node1 in [aut.states[node0.transition[0]], aut.states[node0.transition[1]]]) and (node2 in [aut.states[node0.transition[0]], aut.states[node0.transition[1]]])
    assert node0.transition == node1.transition
    assert len(node1.transition) == 2 and len(node2.transition) == 2

# Close the loop
def loaddumpXML_test():
    aut_first = Automaton()
    aut_first.loadXML(REFERENCE_XMLFRAGMENT)
    aut_second = Automaton()
    aut_second.loadXML(aut_first.dumpXML())
    assert aut_first == aut_second
    aut_first.states[0].transition = []
    assert aut_first != aut_second
    # Comparing string outputs seems too fragile to be a meaningful test
    #assert aut_first.dumpXML(pretty=True, idt_level=1) == REFERENCE_XMLFRAGMENT

def node_copy_test():
    node1 = AutomatonState(id=0, state={"nyan":0, "cat":1}, transition=[0])
    node2 = node1.copy()
    assert node2 is not node1
    assert node2.id == 0
    assert node2.state == {"nyan":0, "cat":1}
    assert node2.transition == [0]

class basic_Automaton_test():
    def setUp(self):
        self.empty_aut = Automaton()
        self.singleton = Automaton(states_or_file=[AutomatonState(id=0, state={"x":0}, transition=[0])])
        self.chain_len = 8
        base_state = dict([("x"+str(k),0) for k in range(self.chain_len)])
        self.chain_aut = Automaton()
        for k in range(self.chain_len):
            base_state["x"+str(k)] = 1
            self.chain_aut.addAutState(AutomatonState(id=k, state=base_state, transition=[(k+1)%self.chain_len]))
            base_state["x"+str(k)] = 0

    def tearDown(self):
        pass

    def test_equality(self):
        assert self.empty_aut == self.empty_aut
        assert self.singleton == self.singleton
        assert self.empty_aut != self.singleton

    def test_size(self):
        assert len(self.empty_aut) == 0  # Should give same result as size()
        assert self.empty_aut.size() == 0
        assert self.singleton.size() == 1
        assert len(self.chain_aut) == self.chain_len

    def test_copy(self):
        A = self.singleton.copy()
        assert len(A) == 1
        B = copy.copy(self.singleton)
        assert A == B
        assert A is not B

        C = self.chain_aut.copy()
        assert len(C) == self.chain_len
        assert C == self.chain_aut
        base_state = dict([("x"+str(k),0) for k in range(self.chain_len)])
        base_state["x0"] = 1
        assert C.findAutState(base_state) != -1
        C.findAutState(base_state).state["x0"] = 0
        assert C != self.chain_aut
<|endoftext|>


<commit_before>import json
import os
from glob import glob

def _render(name, help_text, required, max_value, min_value, default, choices):
    res = ['- **%s.**' % name,
           '  %s\n' % help_text]

    for key, val in zip(['required', 'max', 'min', 'default'], [required, max_value, min_value, default]):
        if val:
            res.append('\n    %s: %s' % (key, val))
    if choices:
        res.append('\n    choices:')
        for choice in choices:
            res.append('      - %s' % choice)
    res.append('\n')
    return res

def render(field):
    return _render(field['name'], field['help_text'], field['required'], field['max'], field['min'], field['default'], field['choices'])

def render_key(field):
    help_text = "Key form :doc:`%s`" % field['reference']
    return _render(field['name'], help_text, field['required'], field['max'], field['min'], field['default'], field['choices'])

def render_key_list(field):
    help_text = "Coma separated keys rom :doc:`%s`" % field['reference']
    return _render(field['name'], help_text, field['required'], field['max'], field['min'], field['default'], field['choices'])

renders = {
    'key': render_key,
    'key_list': render_key_list,
    'default': render
}

with open('fields.json') as f:
    fields = json.load(f)

def collect_group(name, profile_fields):
    res = []
    profile_fields.sort(key=lambda x: x['name'])
    res.append(name.capitalize())
    res.append("_" * len(name))

    for field in profile_fields:
        if field['type'] in renders:
            render_function = renders[field['type']]
        else:
            render_function = renders['default']
        res.extend(render_function(field))

    return res

def profile_to_rst(name):
    with open(name) as f:
        profile = json.load(f)

    res = []
    name = os.path.basename(name)
    res.append(name.replace('.json', '.ini'))
    res.append("=" * len(name))

    profiles = {}

    for field in profile:
        field_template = fields[field['type']].copy()
        field_template.update(**field)
        profiles.setdefault(field_template.get("group", "fields"), []).append(field_template)

    for group in ["fields", "stats", "position", "models", "audio"]:
        group_val = profiles.pop(group, None)
        if group_val:
            res.extend(collect_group(group, group_val))

    # all other groups
    for key, val in profiles.items():
        res.extend(collect_group(key, val))

    return "\n".join(res)

if __name__ == '__main__':
    dirs = glob('profile/*.json')

    for x in dirs:
        res = profile_to_rst(x)
        name = os.path.basename(x.replace('json', 'rst'))
        #print name
        with open('docs/source/%s' % name, 'w') as f:
            f.write(res)

<commit_msg>cleanup<commit_after>import json
import os
from glob import glob

def _render(name, help_text, required, max_value, min_value, default, choices):
    res = ['- **%s.**' % name,
           '  %s\n' % help_text]

    for key, val in zip(['required', 'max', 'min', 'default'], [required, max_value, min_value, default]):
        if val:
            res.append('\n    %s: %s' % (key, val))
    if choices:
        res.append('\n    choices:')
        for choice in choices:
            res.append('      - %s' % choice)
    res.append('\n')
    return res

def render(field):
    return _render(field['name'],
                   field['help_text'],
                   field['required'],
                   field['max'],
                   field['min'],
                   field['default'],
                   field['choices'])

def render_key(field):
    help_text = "Key form :doc:`%s`" % field['reference']
    return _render(field['name'],
                   help_text,
                   field['required'],
                   field['max'],
                   field['min'],
                   field['default'],
                   field['choices'])

def render_key_list(field):
    help_text = "Coma separated keys rom :doc:`%s`" % field['reference']
    return _render(field['name'],
                   help_text, field['required'],
                   field['max'],
                   field['min'],
                   field['default'],
                   field['choices'])

renders = {
    'key': render_key,
    'key_list': render_key_list,
    'default': render
}

with open('fields.json') as f:
    fields = json.load(f)

def collect_group(name, profile_fields):
    res = []
    profile_fields.sort(key=lambda x: x['name'])
    res.append(name.capitalize())
    res.append("_" * len(name))

    for field in profile_fields:
        if field['type'] in renders:
            render_function = renders[field['type']]
        else:
            render_function = renders['default']
        res.extend(render_function(field))
    return res

def profile_to_rst(name):
    with open(name) as f:
        profile = json.load(f)

    res = []
    name = os.path.basename(name)
    res.append(name.replace('.json', '.ini'))
    res.append("=" * len(name))

    profiles = {}

    for field in profile:
        field_template = fields[field['type']].copy()
        field_template.update(**field)
        profiles.setdefault(field_template.get("group", "fields"), []).append(field_template)

    for group in ["fields", "stats", "position", "models", "audio"]:
        group_val = profiles.pop(group, None)
        if group_val:
            res.extend(collect_group(group, group_val))

    # all other groups
    for key, val in profiles.items():
        res.extend(collect_group(key, val))
    return "\n".join(res)

if __name__ == '__main__':
    dirs = glob('profile/*.json')
    for x in dirs:
        res = profile_to_rst(x)
        name = os.path.basename(x.replace('json', 'rst'))
        with open('docs/source/%s' % name, 'w') as f:
            f.write(res)

<|endoftext|>


<commit_before># -*- coding: utf-8 -*-
"""start.py - Generate bash scripts for qsub and execute.
Author: Ian Huston
"""
from __future__ import print_function

import os.path
import configuration
import run_config
import harness
import time
import sys
from helpers import ensurepath
from optparse import OptionParser, OptionGroup
import logging
import subprocess
import helpers

@property
def base_qsub_dict():
    qdict = dict(codedir = run_config.CODEDIR,
                 runname = run_config.PROGRAM_NAME,
                 timelimit = run_config.timelimit,
                 qsublogname = run_config.qsublogname,
                 taskmin = run_config.taskmin,
                 taskmax = run_config.taskmax,
                 hold_jid_list = run_config.hold_jid_list,
                 fotemplatefile = run_config.fotemplatefile,
                 fulltemplatefile = run_config.fulltemplatefile,
                 foscriptname = run_config.foscriptname,
                 fullscriptname = run_config.fullscriptname
                 )
    return qdict

def launch_qsub(qsubscript):
    """Submit the job to the queueing system using qsub.

    Return job id of new job.
    """
    qsubcommand = ["qsub", "-terse", qsubscript]
    newprocess = subprocess.Popen(qsubcommand, stdout=subprocess.PIPE, stderr=subprocess.PIPE)

    result = newprocess.stdout.read()
    error_msg = newprocess.stderr.read()

    if error_msg:
        raise Exception(error_msg)

    job_id = result.rstrip()

    return job_id

def write_out_template(templatefile, newfile, textdict):
    """Write the textdict dictionary using the templatefile to newfile."""
    try:
        f = open(templatefile, "r")
        text = f.read()
    except IOError:
        raise
    finally:
        f.close()

    #Ensure directory exists for new file
    try:
        ensurepath(newfile)
    except IOError:
        raise
    try:
        nf = open(newfile, "w")
        nf.write(text%textdict)
    except IOError:
        raise
    finally:
        nf.close()

def main(argv=None):
    """Process command line options, create qsub scripts and start execution."""

    if not argv:
        argv = sys.argv

    #Default dictionary for templates
    template_dict = base_qsub_dict

    #Parse command line options
    parser = OptionParser()

    loggroup = OptionGroup(parser, "Log Options",
                           "These options affect the verbosity of the log files generated.")
    loggroup.add_option("-q", "--quiet",
                  action="store_const", const=logging.FATAL, dest="loglevel",
                  help="only print fatal error messages")
    loggroup.add_option("-v", "--verbose",
                  action="store_const", const=logging.INFO, dest="loglevel",
                  help="print informative messages")
    loggroup.add_option("--debug",
                  action="store_const", const=logging.DEBUG, dest="loglevel",
                  help="print lots of debugging information",
                  default=run_config.LOGLEVEL)
    parser.add_option_group(loggroup)

    cfggroup = OptionGroup(parser, "Simulation configuration Options",
                           "These options affect the options used by the simulation.")
    cfggroup.add_option("--name", action="store", dest="runname",
                        type="string", help="name of run")
    cfggroup.add_option("--timelimit", action="store", dest="timelimit",
                        type="string", help="time for simulation in format hh:mm:ss")
    cfggroup.add_option("--taskmin", action="store", dest="taskmin",
                        type="string", metavar="NUM", help="minimum task number, default: 1")
    cfggroup.add_option("-t", "--taskmax", action="store", dest="taskmax",
                        type="string", metavar="NUM", help="maximum task number, default: 20")
    parser.add_option_group(cfggroup)

    filegrp = OptionGroup(parser, "File options",
                          "These options override the default choice of template and script files.")
    filegrp.add_option("--fotemplate", action="store", dest="fotemplatefile",
                       type="string", help="first order template file")
    filegrp.add_option("--fulltemplate", action="store", dest="fulltemplatefile",
                       type="string", help="full program template file")
    filegrp.add_option("--foscript", action="store", dest="foscriptname",
                       type="string", help="first order script name")
    filegrp.add_option("--fullscript", action="store", dest="fullscriptname",
                       type="string", help="full program script name")
    parser.add_option_group(filegrp)

    (options, args) = parser.parse_args(args=argv[1:])

    if args:
        raise ValueError("No extra command line arguments are allowed!")

    #Update dictionary with options
    template_dict.update(options)

    helpers.startlogging(log, run_config.logfile, options.loglevel)

    #Log options chosen
    log.debug("Generic template dictionary is %s", template_dict)

    #First order script creation
    fo_dict = template_dict.copy()
    fo_dict["runname"] += "-fo"

    if os.path.isfile(fo_dict["fotemplatefile"]):
        write_out_template(fo_dict["fotemplatefile"],fo_dict["foscriptname"], fo_dict)
    else:
        raise IOError("No template file found at %s!" % fo_dict["fotemplatefile"])

    #Launch first order script and get job id
    try:
        fo_jid = launch_qsub(fo_dict["foscriptname"])
        log.info("Submitted first order script with job id %s.", fo_jid)
    except Exception:
        log.error("Error executing script %s", fo_dict["foscriptname"])
        raise

    #Write second order file with job_id from first
    full_dict = template_dict.copy()
    full_dict["hold_jid_list"] = fo_jid

    if os.path.isfile(full_dict["fulltemplatefile"]):
        write_out_template(full_dict["fulltemplatefile"],full_dict["fullscriptname"], full_dict)
    else:
        raise IOError("No template file found at %s!" % full_dict["fulltemplatefile"])

    #Launch full script and get job id
    try:
        full_jid = launch_qsub(full_dict["fullscriptname"])
        log.info("Submitted full script with job id %s.", full_jid)
    except Exception:
        log.error("Error executing script %s", full_dict["fullscriptname"])
        raise

    return 0

if __name__ == "__main__":
    # Start logging
    log=logging.getLogger()
    try:
        sys.exit(main())
    except Exception as e:
        print("Something went wrong!", file=sys.stderr)
        print(e.message, file=sys.stderr)
        sys.exit(1)

    <commit_msg>Change to using a dictionary for base_qsub_dict.<commit_after># -*- coding: utf-8 -*-
"""start.py - Generate bash scripts for qsub and execute.
Author: Ian Huston
"""
from __future__ import print_function

import os.path
import configuration
import run_config
import harness
import time
import sys
from helpers import ensurepath
from optparse import OptionParser, OptionGroup
import logging
import subprocess
import helpers

#Dictionary of qsub configuration values
base_qsub_dict = dict(codedir = run_config.CODEDIR,
                 runname = run_config.PROGRAM_NAME,
                 timelimit = run_config.timelimit,
                 qsublogname = run_config.qsublogname,
                 taskmin = run_config.taskmin,
                 taskmax = run_config.taskmax,
                 hold_jid_list = run_config.hold_jid_list,
                 fotemplatefile = run_config.fotemplatefile,
                 fulltemplatefile = run_config.fulltemplatefile,
                 foscriptname = run_config.foscriptname,
                 fullscriptname = run_config.fullscriptname
                 )

def launch_qsub(qsubscript):
    """Submit the job to the queueing system using qsub.

    Return job id of new job.
    """
    qsubcommand = ["qsub", "-terse", qsubscript]
    newprocess = subprocess.Popen(qsubcommand, stdout=subprocess.PIPE, stderr=subprocess.PIPE)

    result = newprocess.stdout.read()
    error_msg = newprocess.stderr.read()

    if error_msg:
        raise Exception(error_msg)

    job_id = result.rstrip()

    return job_id

def write_out_template(templatefile, newfile, textdict):
    """Write the textdict dictionary using the templatefile to newfile."""
    try:
        f = open(templatefile, "r")
        text = f.read()
    except IOError:
        raise
    finally:
        f.close()

    #Ensure directory exists for new file
    try:
        ensurepath(newfile)
    except IOError:
        raise
    try:
        nf = open(newfile, "w")
        nf.write(text%textdict)
    except IOError:
        raise
    finally:
        nf.close()

def main(argv=None):
    """Process command line options, create qsub scripts and start execution."""

    if not argv:
        argv = sys.argv

    #Default dictionary for templates
    template_dict = base_qsub_dict.copy()

    #Parse command line options
    parser = OptionParser()

    loggroup = OptionGroup(parser, "Log Options",
                           "These options affect the verbosity of the log files generated.")
    loggroup.add_option("-q", "--quiet",
                  action="store_const", const=logging.FATAL, dest="loglevel",
                  help="only print fatal error messages")
    loggroup.add_option("-v", "--verbose",
                  action="store_const", const=logging.INFO, dest="loglevel",
                  help="print informative messages")
    loggroup.add_option("--debug",
                  action="store_const", const=logging.DEBUG, dest="loglevel",
                  help="print lots of debugging information",
                  default=run_config.LOGLEVEL)
    parser.add_option_group(loggroup)

    cfggroup = OptionGroup(parser, "Simulation configuration Options",
                           "These options affect the options used by the simulation.")
    cfggroup.add_option("--name", action="store", dest="runname",
                        type="string", help="name of run")
    cfggroup.add_option("--timelimit", action="store", dest="timelimit",
                        type="string", help="time for simulation in format hh:mm:ss")
    cfggroup.add_option("--taskmin", action="store", dest="taskmin",
                        type="string", metavar="NUM", help="minimum task number, default: 1")
    cfggroup.add_option("-t", "--taskmax", action="store", dest="taskmax",
                        type="string", metavar="NUM", help="maximum task number, default: 20")
    parser.add_option_group(cfggroup)

    filegrp = OptionGroup(parser, "File options",
                          "These options override the default choice of template and script files.")
    filegrp.add_option("--fotemplate", action="store", dest="fotemplatefile",
                       type="string", help="first order template file")
    filegrp.add_option("--fulltemplate", action="store", dest="fulltemplatefile",
                       type="string", help="full program template file")
    filegrp.add_option("--foscript", action="store", dest="foscriptname",
                       type="string", help="first order script name")
    filegrp.add_option("--fullscript", action="store", dest="fullscriptname",
                       type="string", help="full program script name")
    parser.add_option_group(filegrp)

    (options, args) = parser.parse_args(args=argv[1:])

    if args:
        raise ValueError("No extra command line arguments are allowed!")

    #Update dictionary with options
    template_dict.update(options)

    helpers.startlogging(log, run_config.logfile, options.loglevel)

    #Log options chosen
    log.debug("Generic template dictionary is %s", template_dict)

    #First order script creation
    fo_dict = template_dict.copy()
    fo_dict["runname"] += "-fo"

    if os.path.isfile(fo_dict["fotemplatefile"]):
        write_out_template(fo_dict["fotemplatefile"],fo_dict["foscriptname"], fo_dict)
    else:
        raise IOError("No template file found at %s!" % fo_dict["fotemplatefile"])

    #Launch first order script and get job id
    try:
        fo_jid = launch_qsub(fo_dict["foscriptname"])
        log.info("Submitted first order script with job id %s.", fo_jid)
    except Exception:
        log.error("Error executing script %s", fo_dict["foscriptname"])
        raise

    #Write second order file with job_id from first
    full_dict = template_dict.copy()
    full_dict["hold_jid_list"] = fo_jid

    if os.path.isfile(full_dict["fulltemplatefile"]):
        write_out_template(full_dict["fulltemplatefile"],full_dict["fullscriptname"], full_dict)
    else:
        raise IOError("No template file found at %s!" % full_dict["fulltemplatefile"])

    #Launch full script and get job id
    try:
        full_jid = launch_qsub(full_dict["fullscriptname"])
        log.info("Submitted full script with job id %s.", full_jid)
    except Exception:
        log.error("Error executing script %s", full_dict["fullscriptname"])
        raise

    return 0

if __name__ == "__main__":
    # Start logging
    log=logging.getLogger()
    try:
        sys.exit(main())
    except Exception as e:
        print("Something went wrong!", file=sys.stderr)
        print(e.message, file=sys.stderr)
        sys.exit(1)

    <|endoftext|>