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

import inspect
import logging
import os
import pkgutil
import textwrap
import ujson as json
from argparse import ArgumentParser
from argparse import RawDescriptionHelpFormatter
import traceback
import yaml
from pprint import pprint
import arrow

import zmq

import cif.store
from cif.constants import STORE_ADDR
from cifsdk.constants import REMOTE_ADDR, CONFIG_PATH
from cifsdk.exceptions import AuthError, InvalidSearch
from cifsdk.utils import setup_logging, get_argument_parser, setup_signals

MOD_PATH = os.path.dirname(os.path.abspath(inspect.getfile(inspect.currentframe())))

STORE_PATH = os.path.join(MOD_PATH, "store")
RCVTIMEO = 5000
SNDTIMEO = 2000
LINGER = 3
STORE_DEFAULT = 'sqlite'
STORE_PLUGINS = ['cif.store.dummy', 'cif.store.sqlite', 'cif.store.elasticsearch', 'cif.store.rdflib']

class Store(object):
    def __enter__(self):
        return self

    def __exit__(self, type, value, traceback):
        return self

    def __init__(self, context=zmq.Context.instance(), store_type=STORE_DEFAULT, store_address=STORE_ADDR, **kwargs):

        self.logger = logging.getLogger(__name__)
        self.context = context
        self.store_addr = store_address
        self.store = store_type

        self.router = self.context.socket(zmq.ROUTER)

        self._load_plugin(**kwargs)

    def _load_plugin(self, **kwargs):
        # TODO replace with cif.utils.load_plugin
        self.logger.debug('store is: {}'.format(self.store))
        for loader, modname, is_pkg in pkgutil.iter_modules(cif.store.__path__, 'cif.store.'):
            self.logger.debug('testing store plugin: {}'.format(modname))
            if modname == 'cif.store.{}'.format(self.store) or modname == 'cif.store.z{}'.format(self.store):
                self.logger.debug('Loading plugin: {0}'.format(modname))
                self.store = loader.find_module(modname).load_module(modname)
                self.store = self.store.Plugin(**kwargs)

    def start(self):
        self.token_create_admin()

        self.logger.debug('connecting to router: {}'.format(self.store_addr))
        self.router.connect(self.store_addr)

        self.logger.info('connected')

        self.logger.debug('staring loop')
        while True:
            self.logger.debug('waiting on message')
            m = self.router.recv_multipart()
            self.handle_message(m)

    def handle_message(self, m):
        self.logger.debug('message received')

        self.logger.debug(m)

        id, client_id, null, mtype, token, data = m

        self.logger.debug(data)

        if isinstance(data, basestring):
            try:
                data = json.loads(data)
            except ValueError as e:
                self.logger.error(e)
                self.router.send_multipart([id, json.dumps({"status": "failed"})])
                return

        handler = getattr(self, "handle_" + mtype)
        if handler:
            self.logger.debug("mtype: {0}".format(mtype))
            self.logger.debug('running handler: {}'.format(mtype))

            try:
                rv = handler(token, data)
                rv = {"status": "success", "data": rv}
                self.logger.debug('updating last_active')
                ts = arrow.utcnow().format('YYYY-MM-DDTHH:mm:ss.SSSSS')
                ts = '{}Z'.format(ts)
                self.store.token_last_activity_at(token, timestamp=ts)
            except AuthError as e:
                self.logger.error(e)
                rv = {'status': 'failed', 'message': 'unauthorized'}
            except InvalidSearch as e:
                rv = {'status': 'failed', 'message': 'invalid search'}
            except Exception as e:
                self.logger.error(e)
                traceback.print_exc()
                rv = {"status": "failed"}

            self.router.send_multipart([id, client_id, null, mtype, json.dumps(rv).encode('utf-8')])
        else:
            self.logger.error('message type {0} unknown'.format(mtype))
            self.router.send_multipart([id, null, '0'])

    def handle_ping(self, token, data='[]'):
        self.logger.debug('handling ping message')
        return self.store.ping(token)

    def handle_ping_write(self, token, data='[]'):
        self.logger.debug('handling ping write')
        return self.store.token_write(token)

    def handle_indicators_search(self, token, data):
        if self.store.token_read(token):
            self.logger.debug('searching')
            try:
                x = self.store.indicators_search(token, data)
            except Exception as e:
                raise InvalidSearch('invalid search')
            else:
                return x
        else:
            raise AuthError('invalid token')

    # TODO group check
    def handle_indicators_create(self, token, data):
        if self.store.token_write(token):
            return self.store.indicators_create(token, data)
        else:
            raise AuthError('invalid token')

    def handle_tokens_search(self, token, data):
        if self.store.token_admin(token):
            self.logger.debug('tokens_search')
            return self.store.tokens_search(data)
        else:
            raise AuthError('invalid token')

    def handle_tokens_create(self, token, data):
        if self.store.token_admin(token):
            self.logger.debug('tokens_create')
            return self.store.tokens_create(data)
        else:
            raise AuthError('invalid token')

    def handle_tokens_delete(self, token, data):
        if self.store.token_admin(token):
            return self.store.tokens_delete(data)
        else:
            raise AuthError('invalid token')

    def handle_token_write(self, token, data=None):
        return self.store.token_write(token)

    def handle_tokens_edit(self, token, data):
        if self.store.token_admin(token):
            return self.store.token_edit(data)
        else:
            raise AuthError('invalid token')

    def token_create_admin(self):
        self.logger.info('testing for tokens...')
        if not self.store.tokens_admin_exists():
            self.logger.info('admin token does not exist, generating..')
            rv = self.store.tokens_create({
                'username': u'admin',
                'groups': [u'everyone'],
                'read': u'1',
                'write': u'1',
                'admin': u'1'
            })
            self.logger.info('admin token created: {}'.format(rv['token']))
            return rv['token']
        else:
            self.logger.info('admin token exists...')

    def token_create_smrt(self):
        self.logger.info('generating smrt token')
        rv = self.store.tokens_create({
            'username': u'csirtg-smrt',
            'groups': [u'everyone'],
            'write': u'1',
        })
        self.logger.info('admin token created: {}'.format(rv['token']))
        return rv['token']

    def token_create_hunter(self):
        self.logger.info('generating hunter token')
        rv = self.store.tokens_create({
            'username': u'hunter',
            'groups': [u'everyone'],
            'write': u'1',
        })
        self.logger.info('admin token created: {}'.format(rv['token']))
        return rv['token']

def main():
    p = get_argument_parser()
    p = ArgumentParser(
        description=textwrap.dedent('''\
         Env Variables:
            CIF_RUNTIME_PATH
            CIF_STORE_ADDR

        example usage:
            $ cif-store -d
        '''),
        formatter_class=RawDescriptionHelpFormatter,
        prog='cif-store',
        parents=[p]
    )

    p.add_argument("--store-address", help="specify the store address cif-router is listening on[default: %("
                                             "default)s]", default=STORE_ADDR)

    p.add_argument("--store", help="specify a store type {} [default: %(default)s]".format(', '.join(STORE_PLUGINS)),
                   default=STORE_DEFAULT)

    p.add_argument('--store-type')
    p.add_argument('--nodes')

    p.add_argument('--config', help='specify config path [default %(default)s]', default=CONFIG_PATH)

    p.add_argument('--token-create-admin', help='generate an admin token')
    p.add_argument('--token-create-smrt')
    p.add_argument('--token-create-smrt-remote', default=REMOTE_ADDR)
    p.add_argument('--token-create-hunter')

    args = p.parse_args()

    setup_logging(args)
    logger = logging.getLogger(__name__)
    logger.info('loglevel is: {}'.format(logging.getLevelName(logger.getEffectiveLevel())))

    setup_signals(__name__)

    if args.token_create_smrt:
        with Store(store_type=args.store, nodes=args.nodes) as s:
            t = s.token_create_smrt()
            if t:
                data = {
                    'token': str(t),
                }
                with open(args.token_create_smrt, 'w') as f:
                    f.write(yaml.dump(data, default_flow_style=False))

                logger.info('token config generated: {}'.format(args.token_create_smrt))
            else:
                logger.error('token not created')

    if args.token_create_hunter:
        with Store(store_type=args.store, nodes=args.nodes) as s:
            t = s.token_create_hunter()
            if t:
                data = {
                    'hunter_token': str(t),
                }
                with open(args.token_create_hunter, 'w') as f:
                    f.write(yaml.dump(data, default_flow_style=False))

                logger.info('token config generated: {}'.format(args.token_create_hunter))
            else:
                logger.error('token not created')

    if args.token_create_admin:
        with Store(store_type=args.store, nodes=args.nodes) as s:
            t = s.token_create_admin()
            if t:
                data = {
                    'token': str(t),
                }
                with open(args.token_create_admin, 'w') as f:
                    f.write(yaml.dump(data, default_flow_style=False))

                logger.info('token config generated: {}'.format(args.token_create_admin))
            else:
                logger.error('token not created')

if __name__ == "__main__":
    main()
<commit_msg>bugfix<commit_after>#!/usr/bin/env python

import inspect
import logging
import os
import pkgutil
import textwrap
import ujson as json
from argparse import ArgumentParser
from argparse import RawDescriptionHelpFormatter
import traceback
import yaml
from pprint import pprint
import arrow

import zmq

import cif.store
from cif.constants import STORE_ADDR
from cifsdk.constants import REMOTE_ADDR, CONFIG_PATH
from cifsdk.exceptions import AuthError, InvalidSearch
from cifsdk.utils import setup_logging, get_argument_parser, setup_signals

MOD_PATH = os.path.dirname(os.path.abspath(inspect.getfile(inspect.currentframe())))

STORE_PATH = os.path.join(MOD_PATH, "store")
RCVTIMEO = 5000
SNDTIMEO = 2000
LINGER = 3
STORE_DEFAULT = 'sqlite'
STORE_PLUGINS = ['cif.store.dummy', 'cif.store.sqlite', 'cif.store.elasticsearch', 'cif.store.rdflib']

class Store(object):
    def __enter__(self):
        return self

    def __exit__(self, type, value, traceback):
        return self

    def __init__(self, context=zmq.Context.instance(), store_type=STORE_DEFAULT, store_address=STORE_ADDR, **kwargs):

        self.logger = logging.getLogger(__name__)
        self.context = context
        self.store_addr = store_address
        self.store = store_type

        self.router = self.context.socket(zmq.ROUTER)

        self._load_plugin(**kwargs)

    def _load_plugin(self, **kwargs):
        # TODO replace with cif.utils.load_plugin
        self.logger.debug('store is: {}'.format(self.store))
        for loader, modname, is_pkg in pkgutil.iter_modules(cif.store.__path__, 'cif.store.'):
            self.logger.debug('testing store plugin: {}'.format(modname))
            if modname == 'cif.store.{}'.format(self.store) or modname == 'cif.store.z{}'.format(self.store):
                self.logger.debug('Loading plugin: {0}'.format(modname))
                self.store = loader.find_module(modname).load_module(modname)
                self.store = self.store.Plugin(**kwargs)

    def start(self):
        self.token_create_admin()

        self.logger.debug('connecting to router: {}'.format(self.store_addr))
        self.router.connect(self.store_addr)

        self.logger.info('connected')

        self.logger.debug('staring loop')
        while True:
            self.logger.debug('waiting on message')
            m = self.router.recv_multipart()
            self.handle_message(m)

    def handle_message(self, m):
        self.logger.debug('message received')

        self.logger.debug(m)

        id, client_id, null, mtype, token, data = m

        self.logger.debug(data)

        if isinstance(data, basestring):
            try:
                data = json.loads(data)
            except ValueError as e:
                self.logger.error(e)
                self.router.send_multipart([id, json.dumps({"status": "failed"})])
                return

        handler = getattr(self, "handle_" + mtype)
        if handler:
            self.logger.debug("mtype: {0}".format(mtype))
            self.logger.debug('running handler: {}'.format(mtype))

            try:
                rv = handler(token, data)
                rv = {"status": "success", "data": rv}
                self.logger.debug('updating last_active')
                ts = arrow.utcnow().format('YYYY-MM-DDTHH:mm:ss.SSSSS')
                ts = '{}Z'.format(ts)
                self.store.token_last_activity_at(token, timestamp=ts)
            except AuthError as e:
                self.logger.error(e)
                rv = {'status': 'failed', 'message': 'unauthorized'}
            except InvalidSearch as e:
                rv = {'status': 'failed', 'message': 'invalid search'}
            except Exception as e:
                self.logger.error(e)
                traceback.print_exc()
                rv = {"status": "failed"}

            self.router.send_multipart([id, client_id, null, mtype, json.dumps(rv).encode('utf-8')])
        else:
            self.logger.error('message type {0} unknown'.format(mtype))
            self.router.send_multipart([id, null, '0'])

    def handle_ping(self, token, data='[]'):
        self.logger.debug('handling ping message')
        return self.store.ping(token)

    def handle_ping_write(self, token, data='[]'):
        self.logger.debug('handling ping write')
        return self.store.token_write(token)

    def handle_indicators_search(self, token, data):
        if self.store.token_read(token):
            self.logger.debug('searching')
            try:
                x = self.store.indicators_search(token, data)
            except Exception as e:
                raise InvalidSearch('invalid search')
            else:
                return x
        else:
            raise AuthError('invalid token')

    # TODO group check
    def handle_indicators_create(self, token, data):
        if self.store.token_write(token):
            return self.store.indicators_create(token, data)
        else:
            raise AuthError('invalid token')

    def handle_tokens_search(self, token, data):
        if self.store.token_admin(token):
            self.logger.debug('tokens_search')
            return self.store.tokens_search(data)
        else:
            raise AuthError('invalid token')

    def handle_tokens_create(self, token, data):
        if self.store.token_admin(token):
            self.logger.debug('tokens_create')
            return self.store.tokens_create(data)
        else:
            raise AuthError('invalid token')

    def handle_tokens_delete(self, token, data):
        if self.store.token_admin(token):
            return self.store.tokens_delete(data)
        else:
            raise AuthError('invalid token')

    def handle_token_write(self, token, data=None):
        return self.store.token_write(token)

    def handle_tokens_edit(self, token, data):
        if self.store.token_admin(token):
            return self.store.token_edit(data)
        else:
            raise AuthError('invalid token')

    def token_create_admin(self):
        self.logger.info('testing for tokens...')
        if not self.store.tokens_admin_exists():
            self.logger.info('admin token does not exist, generating..')
            rv = self.store.tokens_create({
                'username': u'admin',
                'groups': [u'everyone'],
                'read': u'1',
                'write': u'1',
                'admin': u'1'
            })
            self.logger.info('admin token created: {}'.format(rv['token']))
            return rv['token']
        else:
            self.logger.info('admin token exists...')

    def token_create_smrt(self):
        self.logger.info('generating smrt token')
        rv = self.store.tokens_create({
            'username': u'csirtg-smrt',
            'groups': [u'everyone'],
            'write': u'1',
        })
        self.logger.info('admin token created: {}'.format(rv['token']))
        return rv['token']

    def token_create_hunter(self):
        self.logger.info('generating hunter token')
        rv = self.store.tokens_create({
            'username': u'hunter',
            'groups': [u'everyone'],
            'write': u'1',
        })
        self.logger.info('admin token created: {}'.format(rv['token']))
        return rv['token']

def main():
    p = get_argument_parser()
    p = ArgumentParser(
        description=textwrap.dedent('''\
         Env Variables:
            CIF_RUNTIME_PATH
            CIF_STORE_ADDR

        example usage:
            $ cif-store -d
        '''),
        formatter_class=RawDescriptionHelpFormatter,
        prog='cif-store',
        parents=[p]
    )

    p.add_argument("--store-address", help="specify the store address cif-router is listening on[default: %("
                                             "default)s]", default=STORE_ADDR)

    p.add_argument("--store", help="specify a store type {} [default: %(default)s]".format(', '.join(STORE_PLUGINS)),
                   default=STORE_DEFAULT)

    p.add_argument('--store-type')
    p.add_argument('--nodes')

    p.add_argument('--config', help='specify config path [default %(default)s]', default=CONFIG_PATH)

    p.add_argument('--token-create-admin', help='generate an admin token')
    p.add_argument('--token-create-smrt')
    p.add_argument('--token-create-smrt-remote', default=REMOTE_ADDR)
    p.add_argument('--token-create-hunter')

    p.add_argument('--remote', help='specify remote')

    args = p.parse_args()

    setup_logging(args)
    logger = logging.getLogger(__name__)
    logger.info('loglevel is: {}'.format(logging.getLevelName(logger.getEffectiveLevel())))

    setup_signals(__name__)

    if args.token_create_smrt:
        with Store(store_type=args.store, nodes=args.nodes) as s:
            t = s.token_create_smrt()
            if t:
                data = {
                    'token': str(t),
                }
                if args.remote:
                    data['remote'] = args.remote
                with open(args.token_create_smrt, 'w') as f:
                    f.write(yaml.dump(data, default_flow_style=False))

                logger.info('token config generated: {}'.format(args.token_create_smrt))
            else:
                logger.error('token not created')

    if args.token_create_hunter:
        with Store(store_type=args.store, nodes=args.nodes) as s:
            t = s.token_create_hunter()
            if t:
                data = {
                    'hunter_token': str(t),
                }
                with open(args.token_create_hunter, 'w') as f:
                    f.write(yaml.dump(data, default_flow_style=False))

                logger.info('token config generated: {}'.format(args.token_create_hunter))
            else:
                logger.error('token not created')

    if args.token_create_admin:
        with Store(store_type=args.store, nodes=args.nodes) as s:
            t = s.token_create_admin()
            if t:
                data = {
                    'token': str(t),
                }
                with open(args.token_create_admin, 'w') as f:
                    f.write(yaml.dump(data, default_flow_style=False))

                logger.info('token config generated: {}'.format(args.token_create_admin))
            else:
                logger.error('token not created')

if __name__ == "__main__":
    main()
<|endoftext|>


<commit_before># vim:sw=4:et:ai

import unittest
import gaphor.UML as UML
import gaphor.diagram as diagram

factory = UML.ElementFactory()

class TestActivityNodes(unittest.TestCase):
    def setUp(self):
        self.pkg = factory.create(UML.Package)
        self.dgm = factory.create(UML.Diagram)
        self.dgm.package = self.pkg

    def createItem(self, ui_item_cls, item_cls):
        item = self.dgm.create(ui_item_cls)
        item.package = self.pkg
        item.subject = factory.create(item_cls)
        return item

    def createFlow(self):
        return self.createItem(diagram.FlowItem, UML.ControlFlow)

    def createActionItem(self):
        return self.createItem(diagram.ActionItem, UML.ActivityNode)

    def createDecisionNode(self):
        return self.createItem(diagram.DecisionNodeItem, UML.DecisionNode)

    def connectNodes(self, source, target, flow):
        source.connect_handle(flow.handles[0])
        target.connect_handle(flow.handles[-1])

    def disconnectNodes(self, flow):
        #print '-' * 10
        #source = flow.handles[0].connected_to
        #target = flow.handles[-1].connected_to
        #source.connect_handle(flow.handles[0])
        #target.connect_handle(flow.handles[-1])
        #print '-' * 10
        pass

    def testMergeNode(self):
        """
        Create decision node UI element and add two incoming flows.
        Element subject should change to MergeNode.
        """

        dnode = self.createDecisionNode()

        assert dnode.subject.__class__ == UML.DecisionNode, 'test problem'

        f1 = self.createFlow()
        f2 = self.createFlow()

        a1 = self.createActionItem()
        a2 = self.createActionItem()

        self.connectNodes(a1, dnode, f1)
        self.connectNodes(a2, dnode, f2)

        self.dgm.canvas.update_now()

        self.assertEquals(dnode.subject.__class__, UML.MergeNode)
        self.assertEquals(dnode.combined, False)
        # fixme: check that previous subject is destroyed

        self.disconnectNodes(f2)

        self.assertEquals(dnode.subject.__class__, UML.DecisionNode)
        self.assertEquals(dnode.combined, False)

    def checkCombinedNode(self, node):
        """
        Check combined UI node element.

        Return subject of UI node element and combined UML node.
        """
        self.assertEquals(node.combined, True)
        self.assertEquals(node.subject.__class__, UML.MergeNode)
        self.assertEquals(len(node.subject.outgoing), 1)
        combined_node = node.subject.outgoing[0].target
        self.assertEquals(combined_node.__class__, UML.DecisionNode)
        self.assertEquals(len(combined_node.incoming), 1)
        return node.subject, combined_node

    def testCombinedNodes(self):
        dnode = self.createDecisionNode()
        f1 = self.createFlow()
        f2 = self.createFlow()
        f3 = self.createFlow()
        f4 = self.createFlow()

        a1 = self.createActionItem()
        a2 = self.createActionItem()
        a3 = self.createActionItem()
        a4 = self.createActionItem()

        self.connectNodes(a1, dnode, f1)
        self.connectNodes(a2, dnode, f2)
        self.connectNodes(dnode, a3, f3)
        self.connectNodes(dnode, a4, f4)

        n1, n2 = self.checkCombinedNode(dnode)
        # node has two incoming flows: f1, f2
        self.assertEquals(len(n1.incoming), 2)
        self.assert_(f1.subject in n1.incoming)
        self.assert_(f2.subject in n1.incoming)

        # combined node has two outgoing flows: f3, f4
        self.assertEquals(len(n2.outgoing), 2)
        self.assert_(f3.subject in n2.outgoing)
        self.assert_(f4.subject in n2.outgoing)

        f5 = self.createFlow()
        a5 = self.createActionItem()
        self.connectNodes(dnode, a5, f5)

        n1, n2 = self.checkCombinedNode(dnode)
        # combined node has new outgoing flow: f5
        self.assertEquals(len(n2.outgoing), 3)
        self.assert_(f5.subject in n2.outgoing)

        # disconnect f5 outgoing flow, node still is combined node
        self.disconnectNodes(f5)
        n1, n2 = self.checkCombinedNode(dnode)
        # f5 is removed from combined node
        self.assertEquals(len(n2.outgoing), 2)
        self.assert_(f5.subject not in n2.outgoing)

        self.disconnectNodes(f4)
        # node should not be combined anymore
        self.assertFalse(dnode.combined)
        # f4 was outgoing flow, so UML node class should merge node
        self.assertEqual(dnode.subject.__class__, UML.MergeNode)
        # f3 should be outgoing flow of UML node
        self.assert_(f3.subject in dnode.subject.outgoing)
        # f1 and f2 should be kept as incoming nodes
        self.assert_(f1.subject in dnode.subject.incoming)
        self.assert_(f2.subject in dnode.subject.incoming)

        f4 = self.createFlow()
        self.connectNodes(dnode, a4, f4)
        self.checkCombinedNode(dnode)

        # now, disconnect f1, dnode should be non-combined decision node
        self.disconnectNodes(f1)

        self.assertFalse(dnode.combined)
        self.assertEqual(dnode.subject.__class__, UML.DecisioNode)
        self.assert_(f1.subject not in dnode.subject.incoming)

if __name__ == '__main__':
    unittest.main()
<commit_msg>- fix typos, disconnecting nodes works now<commit_after># vim:sw=4:et:ai

import unittest
import gaphor.UML as UML
import gaphor.diagram as diagram

factory = UML.ElementFactory()

class TestActivityNodes(unittest.TestCase):
    def setUp(self):
        self.pkg = factory.create(UML.Package)
        self.dgm = factory.create(UML.Diagram)
        self.dgm.package = self.pkg

    def createItem(self, ui_item_cls, item_cls):
        item = self.dgm.create(ui_item_cls)
        item.package = self.pkg
        item.subject = factory.create(item_cls)
        return item

    def createFlow(self):
        return self.createItem(diagram.FlowItem, UML.ControlFlow)

    def createActionItem(self):
        return self.createItem(diagram.ActionItem, UML.ActivityNode)

    def createDecisionNode(self):
        return self.createItem(diagram.DecisionNodeItem, UML.DecisionNode)

    def connectNodes(self, source, target, flow):
        source.connect_handle(flow.handles[0])
        target.connect_handle(flow.handles[-1])

    def disconnectNodes(self, flow):
        print '-' * 10
        source = flow.handles[0].connected_to
        target = flow.handles[-1].connected_to
        source.disconnect_handle(flow.handles[0])
        target.disconnect_handle(flow.handles[-1])
        print '-' * 10

    def testMergeNode(self):
        """
        Create decision node UI element and add two incoming flows.
        Element subject should change to MergeNode.
        """

        dnode = self.createDecisionNode()

        assert dnode.subject.__class__ == UML.DecisionNode, 'test problem'

        f1 = self.createFlow()
        f2 = self.createFlow()

        a1 = self.createActionItem()
        a2 = self.createActionItem()

        self.connectNodes(a1, dnode, f1)
        self.connectNodes(a2, dnode, f2)

        self.dgm.canvas.update_now()

        self.assertEquals(dnode.subject.__class__, UML.MergeNode)
        self.assertEquals(dnode.combined, False)
        # fixme: check that previous subject is destroyed

        self.disconnectNodes(f2)

        self.assertEquals(dnode.subject.__class__, UML.DecisionNode)
        self.assertEquals(dnode.combined, False)

    def checkCombinedNode(self, node):
        """
        Check combined UI node element.

        Return subject of UI node element and combined UML node.
        """
        self.assertEquals(node.combined, True)
        self.assertEquals(node.subject.__class__, UML.MergeNode)
        self.assertEquals(len(node.subject.outgoing), 1)
        combined_node = node.subject.outgoing[0].target
        self.assertEquals(combined_node.__class__, UML.DecisionNode)
        self.assertEquals(len(combined_node.incoming), 1)
        return node.subject, combined_node

    def testCombinedNodes(self):
        dnode = self.createDecisionNode()
        f1 = self.createFlow()
        f2 = self.createFlow()
        f3 = self.createFlow()
        f4 = self.createFlow()

        a1 = self.createActionItem()
        a2 = self.createActionItem()
        a3 = self.createActionItem()
        a4 = self.createActionItem()

        self.connectNodes(a1, dnode, f1)
        self.connectNodes(a2, dnode, f2)
        self.connectNodes(dnode, a3, f3)
        self.connectNodes(dnode, a4, f4)

        n1, n2 = self.checkCombinedNode(dnode)
        # node has two incoming flows: f1, f2
        self.assertEquals(len(n1.incoming), 2)
        self.assert_(f1.subject in n1.incoming)
        self.assert_(f2.subject in n1.incoming)

        # combined node has two outgoing flows: f3, f4
        self.assertEquals(len(n2.outgoing), 2)
        self.assert_(f3.subject in n2.outgoing)
        self.assert_(f4.subject in n2.outgoing)

        f5 = self.createFlow()
        a5 = self.createActionItem()
        self.connectNodes(dnode, a5, f5)

        n1, n2 = self.checkCombinedNode(dnode)
        # combined node has new outgoing flow: f5
        self.assertEquals(len(n2.outgoing), 3)
        self.assert_(f5.subject in n2.outgoing)

        # disconnect f5 outgoing flow, node still is combined node
        self.disconnectNodes(f5)
        n1, n2 = self.checkCombinedNode(dnode)
        # f5 is removed from combined node
        self.assertEquals(len(n2.outgoing), 2)
        self.assert_(f5.subject not in n2.outgoing)

        self.disconnectNodes(f4)
        # node should not be combined anymore
        self.assertFalse(dnode.combined)
        # f4 was outgoing flow, so UML node class should merge node
        self.assertEqual(dnode.subject.__class__, UML.MergeNode)
        # f3 should be outgoing flow of UML node
        self.assert_(f3.subject in dnode.subject.outgoing)
        # f1 and f2 should be kept as incoming nodes
        self.assert_(f1.subject in dnode.subject.incoming)
        self.assert_(f2.subject in dnode.subject.incoming)

        f4 = self.createFlow()
        self.connectNodes(dnode, a4, f4)
        self.checkCombinedNode(dnode)

        # now, disconnect f1, dnode should be non-combined decision node
        self.disconnectNodes(f1)

        self.assertFalse(dnode.combined)
        self.assertEqual(dnode.subject.__class__, UML.DecisionNode)
        self.assert_(f1.subject not in dnode.subject.incoming)

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


<commit_before>import re

class Categories(object):
    """
    this class will handle all the functions related to the categories and recipes
    """
    categories = []
    recipes = []

    def __init__(self, category_name=None, category_owner=None):
        """constructor to initialize the global variables"""

        self.category_name = category_name
        self.category_owner = category_owner

    def create_category(self, category_name, category_owner):
        """This will create new and unique category"""
        personal_categories = [owner_list for owner_list in self.categories
                               if category_owner in owner_list]

        similar_cat_list = [searched_cat_name for searched_cat_name in personal_categories
                            if searched_cat_name[0] == category_name]

        regexcategory_name = "[a-zA-Z0-9- .]"

        if re.match(regexcategory_name, category_name):

            if category_name != '' and category_name.strip():
                if self.categories != []:
                    if similar_cat_list == []:
                        self.categories.append([category_name, category_owner, ])
                        return "success"
                    return "catname_uniqueness"
                self.categories.append([category_name, category_owner, ])
                return "success"
            return "null_empty_field"
        return "catname_pattern"

    def view_category(self, category_owner):
        """This will display the categories"""
        personal_list = [owner_list for owner_list in self.categories if category_owner in owner_list]
        return personal_list

    def edit_category(self, current_name, new_name, category_owner):
        """Update the category name"""
        personal_categories = [owner_list for owner_list in self.categories
                               if category_owner in owner_list]

        similar_cat_list = [searched_cat_name for searched_cat_name in personal_categories
                            if searched_cat_name[0] == new_name]
        regexcategory_name = "[a-zA-Z0-9- .]"
        if re.match(regexcategory_name, new_name):
            if new_name != '' and new_name.strip():
                for catList in personal_categories:
                    if current_name in catList:
                        if similar_cat_list == []:
                            category_name_index = personal_categories.index(catList)
                            personal_categories[category_name_index][0] = new_name
                            """Update the category name in the recipes list"""
                            for recList in self.recipes:
                                if current_name in recList:
                                    for index in range(0, len(self.recipes)):
                                        """loop all the indexes with the current"""
                                        self.recipes[index][1] = new_name
                                    return "success"

                            return "success"
                        return "categoryname_uniqueness"
            return "null_empty_field"
        return "categoryname_pattern"

    def delete_category(self, category_name, category_owner):
        """This will delete the categories from user in session"""
        personal_categories = [owner_list for owner_list in self.categories
                               if category_owner in owner_list]

        specific_category_recipes = [recipes for recipes in self.recipes
                                     if category_owner == recipes[2] and category_name == recipes[1]]
        """delete recipes for the deleted category"""
        for recList in self.recipes:
            if category_name in recList:
                for index in range(0, len(specific_category_recipes)):
                    """loop all the indexes with the recipes of the specific category"""
                del specific_category_recipes[index]

        for catList in personal_categories:
            if category_name in catList:
                category_list_index = personal_categories.index(catList)
                del self.categories[category_list_index]
                del personal_categories[category_list_index]
        return personal_categories

    def create_recipe(self, recipe_name, category_name, category_owner):
        """This will add a recipe to the list"""
        specific_category_recipes = [recipes for recipes in self.recipes
                               if category_name in recipes]

        similar_recipe_list = [searched_recipe_name for searched_recipe_name in specific_category_recipes
                            if searched_recipe_name[0] == recipe_name]

        regexrecipe_name = "[a-zA-Z0-9- .]"

        if re.match(regexrecipe_name, recipe_name):

            if recipe_name != '' and recipe_name.strip():
                if self.recipes != []:
                    if similar_recipe_list == []:
                        self.recipes.append([recipe_name, category_name, category_owner,])
                        return "success"
                    return "recipename_uniqueness"
                self.recipes.append([recipe_name,category_name, category_owner])
                return "success"
            return "null_empty_field"
        return "recipename_pattern"

    def view_recipes(self, category_name, recipe_owner):
        """Displays the specific category's recipes"""
        specific_category_recipes = [recipes for recipes in self.recipes
                                     if recipe_owner == recipes[2] and category_name == recipes[1]]
        return specific_category_recipes

    def edit_recipe(self, current_recipe_name, new_recipe_name, category_name, recipe_owner):
        """Update the category name"""
        specific_category_recipes = [recipes for recipes in self.recipes
                                     if recipe_owner == recipes[2] and category_name == recipes[1]]

        similar_recipe_list = [searched_recipe_name for searched_recipe_name in specific_category_recipes
                            if searched_recipe_name[0] == new_recipe_name]

        regexrecipe_name = "[a-zA-Z0-9- .]"
        if re.match(regexrecipe_name, new_recipe_name):
            if new_recipe_name != '' and new_recipe_name.strip():
                for recList in specific_category_recipes:
                    if current_recipe_name in recList:
                        if similar_recipe_list == []:
                            recipe_name_index = specific_category_recipes.index(recList)
                            specific_category_recipes[recipe_name_index][0] = new_recipe_name
                            return "success"
                        return "recipename_uniqueness"
            return "null_empty_field"
        return "recipename_pattern"

    def delete_recipe(self, recipe_name, category_name, recipe_owner):
        """This will delete the recipes"""
        specific_category_recipes = [recipes for recipes in self.recipes
                                     if recipe_owner == recipes[2] and category_name == recipes[1]]
        for recList in specific_category_recipes:
            if recipe_name in recList:
                recipe_list_index = specific_category_recipes.index(recList)
                del self.recipes[recipe_list_index]
                del specific_category_recipes[recipe_list_index]
        return specific_category_recipes<commit_msg>[Fix] replace the space with underscore in the category name<commit_after>import re

class Categories(object):
    """
    this class will handle all the functions related to the categories and recipes
    """
    categories = []
    recipes = []

    def __init__(self, category_name=None, category_owner=None):
        """constructor to initialize the global variables"""

        self.category_name = category_name
        self.category_owner = category_owner

    def create_category(self, category_name, category_owner):
        """This will create new and unique category"""
        personal_categories = [owner_list for owner_list in self.categories
                               if category_owner in owner_list]

        similar_cat_list = [searched_cat_name for searched_cat_name in personal_categories
                            if searched_cat_name[0] == category_name]

        regexcategory_name = "[a-zA-Z0-9- .]"

        category_name = category_name.replace(" ","_")
        if re.match(regexcategory_name, category_name):

            if category_name != '' and category_name.strip():
                if self.categories != []:
                    if similar_cat_list == []:
                        self.categories.append([category_name, category_owner, ])
                        return "success"
                    return "catname_uniqueness"
                self.categories.append([category_name, category_owner, ])
                return "success"
            return "null_empty_field"
        return "catname_pattern"

    def view_category(self, category_owner):
        """This will display the categories"""
        personal_list = [owner_list for owner_list in self.categories if category_owner in owner_list]
        return personal_list

    def edit_category(self, current_name, new_name, category_owner):
        """Update the category name"""
        personal_categories = [owner_list for owner_list in self.categories
                               if category_owner in owner_list]

        similar_cat_list = [searched_cat_name for searched_cat_name in personal_categories
                            if searched_cat_name[0] == new_name]
        regexcategory_name = "[a-zA-Z0-9- .]"
        if re.match(regexcategory_name, new_name):
            if new_name != '' and new_name.strip():
                for catList in personal_categories:
                    if current_name in catList:
                        if similar_cat_list == []:
                            category_name_index = personal_categories.index(catList)
                            personal_categories[category_name_index][0] = new_name
                            """Update the category name in the recipes list"""
                            for recList in self.recipes:
                                if current_name in recList:
                                    for index in range(0, len(self.recipes)):
                                        """loop all the indexes with the current"""
                                        self.recipes[index][1] = new_name
                                    return "success"

                            return "success"
                        return "categoryname_uniqueness"
            return "null_empty_field"
        return "categoryname_pattern"

    def delete_category(self, category_name, category_owner):
        """This will delete the categories from user in session"""
        personal_categories = [owner_list for owner_list in self.categories
                               if category_owner in owner_list]

        specific_category_recipes = [recipes for recipes in self.recipes
                                     if category_owner == recipes[2] and category_name == recipes[1]]
        """delete recipes for the deleted category"""
        for recList in self.recipes:
            if category_name in recList:
                for index in range(0, len(specific_category_recipes)):
                    """loop all the indexes with the recipes of the specific category"""
                del specific_category_recipes[index]

        for catList in personal_categories:
            if category_name in catList:
                category_list_index = personal_categories.index(catList)
                del self.categories[category_list_index]
                del personal_categories[category_list_index]
        return personal_categories

    def create_recipe(self, recipe_name, category_name, category_owner):
        """This will add a recipe to the list"""
        specific_category_recipes = [recipes for recipes in self.recipes
                               if category_name in recipes]

        similar_recipe_list = [searched_recipe_name for searched_recipe_name in specific_category_recipes
                            if searched_recipe_name[0] == recipe_name]

        regexrecipe_name = "[a-zA-Z0-9- .]"

        if re.match(regexrecipe_name, recipe_name):

            if recipe_name != '' and recipe_name.strip():
                if self.recipes != []:
                    if similar_recipe_list == []:
                        self.recipes.append([recipe_name, category_name, category_owner,])
                        return "success"
                    return "recipename_uniqueness"
                self.recipes.append([recipe_name,category_name, category_owner])
                return "success"
            return "null_empty_field"
        return "recipename_pattern"

    def view_recipes(self, category_name, recipe_owner):
        """Displays the specific category's recipes"""
        specific_category_recipes = [recipes for recipes in self.recipes
                                     if recipe_owner == recipes[2] and category_name == recipes[1]]
        return specific_category_recipes

    def edit_recipe(self, current_recipe_name, new_recipe_name, category_name, recipe_owner):
        """Update the category name"""
        specific_category_recipes = [recipes for recipes in self.recipes
                                     if recipe_owner == recipes[2] and category_name == recipes[1]]

        similar_recipe_list = [searched_recipe_name for searched_recipe_name in specific_category_recipes
                            if searched_recipe_name[0] == new_recipe_name]

        regexrecipe_name = "[a-zA-Z0-9- .]"
        if re.match(regexrecipe_name, new_recipe_name):
            if new_recipe_name != '' and new_recipe_name.strip():
                for recList in specific_category_recipes:
                    if current_recipe_name in recList:
                        if similar_recipe_list == []:
                            recipe_name_index = specific_category_recipes.index(recList)
                            specific_category_recipes[recipe_name_index][0] = new_recipe_name
                            return "success"
                        return "recipename_uniqueness"
            return "null_empty_field"
        return "recipename_pattern"

    def delete_recipe(self, recipe_name, category_name, recipe_owner):
        """This will delete the recipes"""
        specific_category_recipes = [recipes for recipes in self.recipes
                                     if recipe_owner == recipes[2] and category_name == recipes[1]]
        for recList in specific_category_recipes:
            if recipe_name in recList:
                recipe_list_index = specific_category_recipes.index(recList)
                del self.recipes[recipe_list_index]
                del specific_category_recipes[recipe_list_index]
        return specific_category_recipes<|endoftext|>


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

# 1st TODO: FileData
# 2nd TODO: unify Input, Output, and Computed: api for sources, splits, etc.
# 3rd TODO: make it all work

# TODO: right now we assume that input files are pre-split.

# Copyright 2008 Brigham Young University
#
# This file is part of Mrs.
#
# Mrs is free software: you can redistribute it and/or modify it under the
# terms of the GNU General Public License as published by the Free Software
# Foundation, either version 3 of the License, or (at your option) any later
# version.
#
# Mrs is distributed in the hope that it will be useful, but WITHOUT ANY
# WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS
# FOR A PARTICULAR PURPOSE.  See the GNU General Public License for more
# details.
#
# You should have received a copy of the GNU General Public License along with
# Mrs.  If not, see <http://www.gnu.org/licenses/>.
#
# Inquiries regarding any further use of the Materials contained on this site,
# please contact the Copyright Licensing Office, Brigham Young University,
# 3760 HBLL, Provo, UT 84602, (801) 422-9339 or 422-3821, e-mail
# copyright@byu.edu.

import threading, os
from heapq import heappush

from io import fileformat, openreader, HexFormat

# TODO: cache data to disk when memory usage is high
class Bucket(object):
    """Hold data from a source or for a split.

    Data can be manually dumped to disk, in which case the data will be saved
    to the given filename with the specified format.  Eventually Bucket will
    be upgraded to automatically cache data to disk if they get too large to
    stay in memory.

    >>> b = Bucket()
    >>> b.append((4, 'test'))
    >>> b.collect([(3, 'a'), (1, 'This'), (2, 'is')])
    >>> ' '.join([value for key, value in b])
    'test a This is'
    >>> b.sort = True
    >>> ' '.join([value for key, value in b])
    'This is a test'
    >>>
    """
    def __init__(self, sort=False, filename=None, format=HexFormat):
        self._data = []
        self.sort = sort
        self.filename = filename

    def append(self, x):
        """Collect a single key-value pair
        """
        if self.sort:
            heappush(self._data, x)
        else:
            self._data.append(x)

    def collect(self, pairiter):
        """Collect all key-value pairs from the given iterable

        The collection can be a generator or a Mrs format.  This will block if
        the iterator blocks.
        """
        data = self._data
        if self.sort:
            for kvpair in pairiter:
                heappush(data, kvpair)
                data.append(kvpair)
        else:
            for kvpair in pairiter:
                data.append(kvpair)

    def __iter__(self):
        if self.sort:
            return iter(sorted(self._data))
        else:
            return iter(self._data)

    # TODO: write doctest for dump
    def dump(self):
        assert(self.filename is not None)
        f = openfile(self.filename, 'w')
        for key, value in self:
            f.write(key, value)
        f.close()

class DataSet(object):
    """Manage input to or output from a map or reduce operation.
    """
    def __init__(self, directory=None, format=HexFormat):
        if directory is None:
            from tempfile import mkdtemp
            self.directory = mkdtemp()
            self.temp = True
        else:
            self.directory = directory
            self.temp = False

    #def __len__(self):
    #    """Number of splits in this DataSet."""
    #    return 0

    #def __getitem__(self, split):
    #    """Retrieve a seq of URLs for a particular split in this DataSet."""
    #    raise IndexError

    #def ready(self, split=None):
    #    """Whether or not a split is ready to be retrieved.

    #    If the split is not specified, ready() tells whether or not
    #    all splits are ready.
    #    """
    #    return True

class Output(DataSet):
    """Collect output from a map or reduce task."""
    def __init__(self, partition, n, **kwds):
        super(Output, self).__init__(**kwds)

        self.partition = partition
        self.n = n
        self.splits = [Bucket(format=format, filename=self.path(i))
                for i in xrange(n)]

    def close(self):
        if self.temp:
            os.removedirs(self.directory)

    def collect(self, itr):
        """Collect all of the key-value pairs from the given iterator."""
        n = self.n
        if n == 1:
            bucket = self.splits[0]
            bucket.collect(itr)
        else:
            partition = self.partition
            for kvpair in itr:
                key, value = kvpair
                split = partition(key, n)
                bucket = self.splits[split]
                bucket.append(kvpair)

    def path(self, index):
        """Return the path to the output split for the given index."""
        filename = "split_%s.%s" % (index, self.format)
        return os.path.join(self.directory, filename)

    def savetodisk(self):
        """Write out all of the key-value pairs to files."""
        for bucket in self.splits:
            bucket.dump()

class FileData(DataSet):
    """A list of static files or urls to be used as input to an operation.

    #>>> urls = ['http://aml.cs.byu.edu/', 'http://www.cs.byu.edu/']
    #>>> data = FileData(urls)
    #>>> len(data)
    #2
    #>>> data[1]
    #'http://www.cs.byu.edu/'
    #>>>
    """
    def __init__(self, urls):
        self._urls = tuple(urls)

    def fetchall(self):
        """Download all of the files
        """
        # TODO: set a maximum number of files to read at the same time (do we
        # really want to have 500 sockets open at once?)

        # TODO: setup a bucket for each URL, setup io.net.download for
        # each url, and run a twisted reactor loop to download it all
        for url in self._urls:
            reader = openreader(url)
            buf.deferred.addCallback(self.callback, reader)

    def callback(self, value, reader):
        pass

    #def __len__(self):
    #    """How many urls are in the DataSet."""
    #    return len(self._urls)

    #def __getitem__(self, i):
    #    """Retrieve a sequence containing the URL for the specified file."""
    #    return (self._urls[i],)

class ComputedData(DataSet):
    """Manage input to or output from a map or reduce operation.

    The data are evaluated lazily.  A DataSet knows how to generate or
    regenerate its contents.  It can also decide whether to save the data to
    permanent storage or to leave them in memory on the slaves.
    """
    def __init__(self, input, func, nparts, outdir, parter=None,
            registry=None):
        self.input = input
        self.ntasks = len(self.input)
        self.outdir = outdir
        if registry is None:
            from registry import Registry
            self.registry = Registry()
        else:
            self.registry = registry
        self.nparts = nparts

        self.func_name = self.registry.as_name(func)
        if parter is None:
            self.part_name = ''
        else:
            self.part_name = self.registry.as_name(parter)

        self.tasks_made = False
        self.tasks_todo = []
        self.tasks_done = []
        self.tasks_active = []

        # TODO: store a mapping from tasks to hosts and a map from hosts to
        # tasks.  This way you can know where to find data.  You also know
        # which hosts to restart in case of failure.

    def __len__(self):
        return self.nparts

    def __getitem__(self, split):
        if not self.ready():
            raise IndexError
        else:
            return [task.outurls[split] for task in self.tasks_done
                    if task.outurls[split] != '']

    def ready(self):
        if self.tasks_made and not self.tasks_todo and not self.tasks_active:
            return True
        else:
            return False

    def get_task(self):
        """Return the next available task"""
        if self.tasks_todo:
            task = self.tasks_todo.pop()
            return task
        else:
            return

    def print_status(self):
        active = len(self.tasks_active)
        todo = len(self.tasks_todo)
        done = len(self.tasks_done)
        total = active + todo + done
        print 'Completed: %s/%s, Active: %s' % (done, total, active)

class MapData(ComputedData):
    def __init__(self, input, mapper, nparts, outdir, parter=None,
            registry=None):
        ComputedData.__init__(self, input, mapper, nparts, outdir,
                parter=None, registry=registry)

    def make_tasks(self):
        from mapreduce import MapTask
        for taskid in xrange(self.ntasks):
            task = MapTask(taskid, self.input, self.registry, self.func_name,
                    self.part_name, self.outdir, self.nparts)
            task.dataset = self
            self.tasks_todo.append(task)
        self.tasks_made = True

    def run_serial(self):
        pass
        input_files = [io.openfile(filename) for filename in self.inputs]
        all_input = chain(*input_files)
        map_output = mrs_map(registry['mapper'], all_input)

class ReduceData(ComputedData):
    def __init__(self, input, reducer, nparts, outdir, parter=None,
            registry=None):
        ComputedData.__init__(self, input, reducer, nparts, outdir,
                parter=None, registry=registry)

    def make_tasks(self):
        from mapreduce import ReduceTask
        for taskid in xrange(self.ntasks):
            task = ReduceTask(taskid, self.input, self.registry,
                    self.func_name, self.outdir)
            task.dataset = self
            self.tasks_todo.append(task)
        self.tasks_made = True

def test_datasets():
    import doctest
    doctest.testmod()

if __name__ == "__main__":
    test_datasets()

# vim: et sw=4 sts=4
<commit_msg>got slicing working for low-level DataSet<commit_after>#!/usr/bin/env python

# 1st TODO: FileData
# 2nd TODO: unify Input, Output, and Computed: api for sources, splits, etc.
# 3rd TODO: make it all work

# TODO: right now we assume that input files are pre-split.

# Copyright 2008 Brigham Young University
#
# This file is part of Mrs.
#
# Mrs is free software: you can redistribute it and/or modify it under the
# terms of the GNU General Public License as published by the Free Software
# Foundation, either version 3 of the License, or (at your option) any later
# version.
#
# Mrs is distributed in the hope that it will be useful, but WITHOUT ANY
# WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS
# FOR A PARTICULAR PURPOSE.  See the GNU General Public License for more
# details.
#
# You should have received a copy of the GNU General Public License along with
# Mrs.  If not, see <http://www.gnu.org/licenses/>.
#
# Inquiries regarding any further use of the Materials contained on this site,
# please contact the Copyright Licensing Office, Brigham Young University,
# 3760 HBLL, Provo, UT 84602, (801) 422-9339 or 422-3821, e-mail
# copyright@byu.edu.

import threading, os
from heapq import heappush

from io import fileformat, openreader, HexFormat

# TODO: cache data to disk when memory usage is high
class Bucket(object):
    """Hold data from a source or for a split.

    Data can be manually dumped to disk, in which case the data will be saved
    to the given filename with the specified format.  Eventually Bucket will
    be upgraded to automatically cache data to disk if they get too large to
    stay in memory.

    >>> b = Bucket()
    >>> b.append((4, 'test'))
    >>> b.collect([(3, 'a'), (1, 'This'), (2, 'is')])
    >>> ' '.join([value for key, value in b])
    'test a This is'
    >>> b.sort = True
    >>> ' '.join([value for key, value in b])
    'This is a test'
    >>>
    """
    def __init__(self, sort=False, filename=None, format=HexFormat):
        self._data = []
        self.sort = sort
        self.filename = filename

    def append(self, x):
        """Collect a single key-value pair
        """
        if self.sort:
            heappush(self._data, x)
        else:
            self._data.append(x)

    def collect(self, pairiter):
        """Collect all key-value pairs from the given iterable

        The collection can be a generator or a Mrs format.  This will block if
        the iterator blocks.
        """
        data = self._data
        if self.sort:
            for kvpair in pairiter:
                heappush(data, kvpair)
                data.append(kvpair)
        else:
            for kvpair in pairiter:
                data.append(kvpair)

    def __iter__(self):
        if self.sort:
            return iter(sorted(self._data))
        else:
            return iter(self._data)

    # TODO: write doctest for dump
    def dump(self):
        assert(self.filename is not None)
        f = openfile(self.filename, 'w')
        for key, value in self:
            f.write(key, value)
        f.close()

class DataSet(object):
    """Manage input to or output from a map or reduce operation.

    A DataSet is naturally a two-dimensional list.  There are some number of
    sources, and for each source, there are one or more splits.

    Low-level Testing.  Normally a DataSet holds Buckets, but for now we'll be
    loose for testing purposes.  This also makes it clear how slicing works,
    so it's not a waste of space.
    >>> ds = DataSet(sources=5, splits=3)
    >>> len(ds)
    5
    >>> ds[0, 0] = 'zero'
    >>> ds[0, 1] = 'one'
    >>> ds[0, 2] = 'two'
    >>> ds[1, 1] = 'hello'
    >>> print ds[0, 1]
    one
    >>> print ds[1, 0]
    None
    >>> print ds[0, 1:]
    ['one', 'two']
    >>> print ds[0, :]
    ['zero', 'one', 'two']
    >>> print ds[0:2, 1:3]
    [['one', 'two'], ['hello', None]]
    >>> print ds[:, 1]
    ['one', 'hello', None, None, None]
    >>>
    """
    def __init__(self, sources=0, splits=0, directory=None, format=HexFormat):
        if directory is None:
            from tempfile import mkdtemp
            self.directory = mkdtemp()
            self.temp = True
        else:
            self.directory = directory
            self.temp = False

        # For now assume that all sources have the same # of splits.
        self._data = [[None] * splits for i in xrange(sources)]

    def __setitem__(self, item, value):
        """Set an item.

        For now, you can't set a split.
        """
        # Separate the two dimensions:
        try:
            part1, part2 = item
        except (TypeError, ValueError):
            raise TypeError("Requires a pair of items.")

        self._data[part1][part2] = value

    def __getitem__(self, item):
        """Retrieve an item or split.

        At the moment, we're not very consistent about whether what we return
        is a view or a [shallow] copy.  Write at your own risk.
        """
        # Separate the two dimensions:
        try:
            part1, part2 = item
        except (TypeError, ValueError):
            raise TypeError("Requires a pair of items.")

        isslice1 = isinstance(part1, slice)
        isslice2 = isinstance(part1, slice)

        data = self._data
        if isslice1:
            lst = []
            wild_goose_chase = True
            for sourcelst in data[part1]:
                try:
                    lst.append(sourcelst[part2])
                    wild_goose_chase = False
                except IndexError:
                    lst.append([])

            if wild_goose_chase:
                # every sourcelst[part2] raised an indexerror
                raise IndexError("No items matching %s" % part2)

            return lst

        else:
            return self._data[part1][part2]

    def __len__(self):
        """Number of sources in this DataSet."""
        return len(self._data)

class Output(DataSet):
    """Collect output from a map or reduce task."""
    def __init__(self, partition, nsplits, **kwds):
        super(Output, self).__init__(**kwds)

        self.partition = partition
        # One source and nsplits splits
        self._data = [[Bucket(format=format, filename=self.path(i))
                for i in xrange(nsplits)]]

    def close(self):
        if self.temp:
            os.removedirs(self.directory)

    def collect(self, itr):
        """Collect all of the key-value pairs from the given iterator."""
        buckets = self[0, :]
        n = self.nsplits
        if n == 1:
            bucket = self.splits[0]
            bucket.collect(itr)
        else:
            partition = self.partition
            nsplits = self.nsplits
            for kvpair in itr:
                key, value = kvpair
                split = partition(key, nsplits)
                bucket = self.splits[split]
                bucket.append(kvpair)

    def path(self, index):
        """Return the path to the output split for the given index."""
        filename = "split_%s.%s" % (index, self.format)
        return os.path.join(self.directory, filename)

    def savetodisk(self):
        """Write out all of the key-value pairs to files."""
        for bucket in self.splits:
            bucket.dump()

class FileData(DataSet):
    """A list of static files or urls to be used as input to an operation.

    #>>> urls = ['http://aml.cs.byu.edu/', 'http://www.cs.byu.edu/']
    #>>> data = FileData(urls)
    #>>> len(data)
    #2
    #>>> data[1]
    #'http://www.cs.byu.edu/'
    #>>>
    """
    def __init__(self, urls):
        self._urls = tuple(urls)

    def fetchall(self):
        """Download all of the files
        """
        # TODO: set a maximum number of files to read at the same time (do we
        # really want to have 500 sockets open at once?)

        # TODO: setup a bucket for each URL, setup io.net.download for
        # each url, and run a twisted reactor loop to download it all
        for url in self._urls:
            reader = openreader(url)
            buf.deferred.addCallback(self.callback, reader)

    def callback(self, value, reader):
        pass

    #def __len__(self):
    #    """How many urls are in the DataSet."""
    #    return len(self._urls)

    #def __getitem__(self, i):
    #    """Retrieve a sequence containing the URL for the specified file."""
    #    return (self._urls[i],)

class ComputedData(DataSet):
    """Manage input to or output from a map or reduce operation.

    The data are evaluated lazily.  A DataSet knows how to generate or
    regenerate its contents.  It can also decide whether to save the data to
    permanent storage or to leave them in memory on the slaves.
    """
    def __init__(self, input, func, nparts, outdir, parter=None,
            registry=None):
        self.input = input
        self.ntasks = len(self.input)
        self.outdir = outdir
        if registry is None:
            from registry import Registry
            self.registry = Registry()
        else:
            self.registry = registry
        self.nparts = nparts

        self.func_name = self.registry.as_name(func)
        if parter is None:
            self.part_name = ''
        else:
            self.part_name = self.registry.as_name(parter)

        self.tasks_made = False
        self.tasks_todo = []
        self.tasks_done = []
        self.tasks_active = []

        # TODO: store a mapping from tasks to hosts and a map from hosts to
        # tasks.  This way you can know where to find data.  You also know
        # which hosts to restart in case of failure.

    def __len__(self):
        return self.nparts

    def __getitem__(self, split):
        if not self.ready():
            raise IndexError
        else:
            return [task.outurls[split] for task in self.tasks_done
                    if task.outurls[split] != '']

    def ready(self):
        if self.tasks_made and not self.tasks_todo and not self.tasks_active:
            return True
        else:
            return False

    def get_task(self):
        """Return the next available task"""
        if self.tasks_todo:
            task = self.tasks_todo.pop()
            return task
        else:
            return

    def print_status(self):
        active = len(self.tasks_active)
        todo = len(self.tasks_todo)
        done = len(self.tasks_done)
        total = active + todo + done
        print 'Completed: %s/%s, Active: %s' % (done, total, active)

class MapData(ComputedData):
    def __init__(self, input, mapper, nparts, outdir, parter=None,
            registry=None):
        ComputedData.__init__(self, input, mapper, nparts, outdir,
                parter=None, registry=registry)

    def make_tasks(self):
        from mapreduce import MapTask
        for taskid in xrange(self.ntasks):
            task = MapTask(taskid, self.input, self.registry, self.func_name,
                    self.part_name, self.outdir, self.nparts)
            task.dataset = self
            self.tasks_todo.append(task)
        self.tasks_made = True

    def run_serial(self):
        pass
        input_files = [io.openfile(filename) for filename in self.inputs]
        all_input = chain(*input_files)
        map_output = mrs_map(registry['mapper'], all_input)

class ReduceData(ComputedData):
    def __init__(self, input, reducer, nparts, outdir, parter=None,
            registry=None):
        ComputedData.__init__(self, input, reducer, nparts, outdir,
                parter=None, registry=registry)

    def make_tasks(self):
        from mapreduce import ReduceTask
        for taskid in xrange(self.ntasks):
            task = ReduceTask(taskid, self.input, self.registry,
                    self.func_name, self.outdir)
            task.dataset = self
            self.tasks_todo.append(task)
        self.tasks_made = True

def test_datasets():
    import doctest
    doctest.testmod()

if __name__ == "__main__":
    test_datasets()

# vim: et sw=4 sts=4
<|endoftext|>


<commit_before>"""Make modules of observations based on cooccurence networks and collapse table"""

import os
import shutil
import glob
import sys

import networkx as nx
import numpy as np
import pysurvey as ps

import general
from biom import load_table, Table
from scipy.stats import spearmanr, pearsonr
from pysurvey import SparCC as sparcc
from functools import partial

def paired_correlations_from_table(table, correl_method=spearmanr, p_adjust=general.bh_adjust):
    """Takes a biom table and finds correlations between all pairs of otus."""
    correls = list()

    for (data_i, otu_i, metadata_i), (data_j, otu_j, metadata_j) in table.iter_pairwise(axis='observation'):
        correl = correl_method(data_i, data_j)
        correls.append([str(otu_i), str(otu_j), correl[0], correl[1]])

    # adjust p-value if desired
    if p_adjust is not None:
        p_adjusted = p_adjust([i[3] for i in correls])
        for i in xrange(len(correls)):
            correls[i].append(p_adjusted[i])

    header = ['feature1', 'feature2', 'r', 'p']
    if p_adjust is not None:
        header.append('adjusted_p')

    return correls, header

def boostrapped_correlation_multi(tup, temp_folder, cor_temp):
    """tup is a tuple where tup[0] is num and tup[1] is the file name"""
    df = ps.read_txt(tup[1], verbose=False)
    cor = ps.basis_corr(df, oprint=False)[0]
    ps.write_txt(cor, temp_folder+cor_temp+str(tup[0])+".txt", T=False)

def sparcc_correlations_multi(table, p_adjust=general.bh_adjust, temp_folder=os.getcwd()+"/temp/",
                              boot_temp="bootstrap_", cor_temp="cor_", table_temp="temp_table.txt",
                              bootstraps=100):
    """Calculate correlations with sparcc"""
    import multiprocessing

    os.mkdir(temp_folder)
    pool = multiprocessing.Pool(multiprocessing.cpu_count()-1)

    # make tab delimited, delete first line and reread in
    with open(temp_folder+table_temp, 'w') as f:
        f.write('\n'.join(table.to_tsv().split("\n")[1:]))
    df = ps.read_txt(temp_folder+table_temp, verbose=False)

    sparcc.make_bootstraps(df, bootstraps, boot_temp+"#.txt", temp_folder)
    cor, cov = ps.basis_corr(df, oprint=False)

    pfun = partial(boostrapped_correlation_multi, temp_folder=temp_folder, cor_temp=cor_temp)
    tups = enumerate(glob.glob(temp_folder+boot_temp+"*.txt"))
    pool.map(pfun, tups)
    pool.close()
    pool.join()

    p_vals = sparcc.get_pvalues(cor, temp_folder+cor_temp+"#.txt", bootstraps)
    # generate correls
    correls = list()
    for i in xrange(len(cor.index)):
        for j in xrange(i+1, len(cor.index)):
            correls.append([str(cor.index[i]), str(cor.index[j]), cor.iat[i, j], p_vals.iat[i, j]])

    # adjust p-value if desired
    if p_adjust is not None:
        p_adjusted = p_adjust([i[3] for i in correls])
        for i in xrange(len(correls)):
            correls[i].append(p_adjusted[i])

    header = ['feature1', 'feature2', 'r', 'p']
    if p_adjust is not None:
        header.append('adjusted_p')

    # cleanup, remove all of bootstraps folder
    shutil.rmtree(temp_folder)

    return correls, header

def boostrapped_correlation(in_file, temp_folder, cor_temp, num):
    df = ps.read_txt(in_file, verbose=False)
    cor = ps.basis_corr(df, oprint=False)[0]
    ps.write_txt(cor, temp_folder+cor_temp+str(num)+".txt", T=False)

def sparcc_correlations(table, p_adjust=general.bh_adjust):
    """"""
    # setup
    temp_folder = "temp3/"
    boot_temp = "bootstrap_"
    cor_temp = "cor_"
    table_temp = "temp_table.txt"
    os.mkdir(temp_folder)
    bootstraps = 100

    # make tab delimited, delete first line and reread in
    with open(temp_folder+table_temp, 'w') as f:
        f.write('\n'.join(table.to_tsv().split("\n")[1:]))
    df = ps.read_txt(temp_folder+table_temp, verbose=False)

    # calculate correlations
    cor, cov = ps.basis_corr(df, oprint=False)

    # calculate p-values
    sparcc.make_bootstraps(df, bootstraps, boot_temp+"#.txt", temp_folder)
    for i, file in enumerate(glob.glob(temp_folder+boot_temp+"*.txt")):
        boostrapped_correlation(file, temp_folder, cor_temp, i)

    p_vals = sparcc.get_pvalues(cor, temp_folder+cor_temp+"#.txt", bootstraps)
    # generate correls
    correls = list()
    for i in xrange(len(cor.index)):
        for j in xrange(i+1, len(cor.index)):
            correls.append([str(cor.index[i]), str(cor.index[j]), cor.iat[i, j], p_vals.iat[i, j]])

    # adjust p-value if desired
    if p_adjust is not None:
        p_adjusted = p_adjust([i[3] for i in correls])
        for i in xrange(len(correls)):
            correls[i].append(p_adjusted[i])

    header = ['feature1', 'feature2', 'r', 'p']
    if p_adjust is not None:
        header.append('adjusted_p')

    # cleanup, remove all of bootstraps folder
    shutil.rmtree(temp_folder)

    return correls, header

def paired_correlations_from_table_with_outlier_removal(table, good_samples, min_keep=10, correl_method=spearmanr,
                                                        p_adjust=general.bh_adjust):
    """Takes a biom table and finds correlations between all pairs of otus."""
    correls = list()

    for (data_i, otu_i, metadata_i), (data_j, otu_j, metadata_j) in table.iter_pairwise(axis='observation'):
        samp_union = np.union1d(good_samples[otu_i], good_samples[otu_j])
        # remove zero zero points
        # samp_union = [ind for i, ind in enumerate(samp_union) if data_i[i]!=0 and data_j[i]!=0]
        if len(samp_union) > min_keep:
            correl = correl_method(data_i[samp_union], data_j[samp_union])
            correls.append([str(otu_i), str(otu_j), correl[0], correl[1]])

    # adjust p-value if desired
    if p_adjust is not None:
        p_adjusted = p_adjust([i[3] for i in correls])
        for i in xrange(len(correls)):
            correls[i].append(p_adjusted[i])

    header = ['feature1', 'feature2', 'r', 'p']
    if p_adjust is not None:
        header.append('adjusted_p')

    return correls, header

def make_modules(graph, k=3):
    """make modules with networkx k-clique communities and annotate network
    :type graph: networkx graph
    """
    cliques = list(nx.k_clique_communities(graph, k))
    cliques = [list(i) for i in cliques]
    for i, clique in enumerate(cliques):
        for node in clique:
            graph.node[node]['k_' + str(k)] = i
    return graph, cliques

def make_modules_multik(graph, k=None):
    """make modules with networkx k-clique communities and annotate network"""
    if k is None:
        k = [2, 3, 4, 5, 6]
    communities = dict()
    for k_val in list(k):
        cliques = list(nx.k_clique_communities(graph, k_val))
        cliques = [list(i) for i in cliques]
        communities[k_val] = cliques
        for i, clique in enumerate(cliques):
            for node in clique:
                graph.node[node]['k_' + str(k_val)] = i
    return graph, {k: list(v) for k, v in communities.iteritems()}

def collapse_modules(table, cliques, prefix="module_"):
    """collapse created modules in a biom table"""
    in_module = set()
    modules = np.zeros((len(cliques), table.shape[1]))

    # for each clique merge values and print cliques to file
    with open("cliques.txt", 'w') as f:
        for i, clique in enumerate(cliques):
            in_module = in_module | set(clique)
            f.write(prefix+str(i)+'\t'+','.join([str(i) for i in clique])+'\n')
            for feature in clique:
                try:
                    modules[i] += table.data(feature, axis='observation')
                except:
                    print table.ids(axis="observation")
                    print len(table.ids(axis="observation"))
                    print table.ids(axis="observation")[0]
                    print feature
                    print type(feature)
                    print unicode(feature)
                    print table.shape
                    sys.exit()
    table.filter(in_module, axis='observation')

    # make new table
    new_table_matrix = np.concatenate((table.matrix_data.toarray(), modules))
    new_table_obs = list(table.ids(axis='observation')) + [prefix + str(i) for i in range(0, len(cliques))]
    return Table(new_table_matrix, new_table_obs, table.ids())

def collapse_modules_multik(table, cliques, prefix="module_"):
    """collapse created modules in a biom table"""
    in_module = set()
    modules = np.zeros((len(cliques), table.shape[1]))

    # for each clique merge values
    for i, clique in enumerate(cliques):
        in_module = in_module | set(clique)
        for feature in clique:
            modules[i] += table.data(feature, axis='observation')
    table.filter(in_module, axis='observation')

    # make new table
    new_table_matrix = np.concatenate((table.matrix_data.toarray(), modules))
    new_table_obs = list(table.ids(axis='observation')) + [prefix + str(i) for i in range(0, len(cliques))]
    return Table(new_table_matrix, new_table_obs, table.ids())

def module_maker(args):
    # correlation and p-value adjustment methods
    correl_methods = {'spearman': spearmanr, 'pearson': pearsonr, 'sparcc': sparcc_correlations_multi}
    p_methods = {'bh': general.bh_adjust, 'bon': general.bonferroni_adjust}
    correl_method = correl_methods[args.correl_method]
    if args.p_adjust is not None:
        p_adjust = p_methods[args.p_adjust]
    else:
        p_adjust = None

    # get features to be correlated and extract metadata
    table = load_table(args.input)
    metadata = general.get_metadata_from_table(table)
    print "Table loaded: " + str(table.shape[0]) + " observations"
    print ""

    # make new output directory and change to it
    if args.output is not None:
        os.makedirs(args.output)
        os.chdir(args.output)

    # convert to relative abundance and filter
    if args.min_sample != None:
        table_filt = general.filter_table(table, args.min_sample)
        print "Table filtered: " + str(table_filt.shape[0]) + " observations"
        print ""
    else:
        table_filt = table

    # correlate feature
    if correl_method in [spearmanr, pearsonr]:
        if args.outlier_removal:
            # remove outlier observations
            # first attempt with just looking at individual otu's
            good_samples = general.remove_outliers(table_filt)
            print "Outliers removed: " + str(len(good_samples)) + " observations"
            print ""
            correls, correl_header = paired_correlations_from_table_with_outlier_removal(table_filt, good_samples,
                                                                                correl_method, p_adjust)
        else:
            # correlate feature
            correls, correl_header = paired_correlations_from_table(table_filt, correl_method, p_adjust)
    else:
        correls, correl_header = sparcc_correlations_multi(table_filt, p_adjust)
    general.print_delimited('correls.txt', correls, correl_header)

    print "Features Correlated"

    # make correlation network
    net = general.correls_to_net(correls, conet=True, metadata=metadata, min_p=args.min_p)
    print "Network Generated"
    print "number of nodes: " + str(net.number_of_nodes())
    print "number of edges: " + str(net.number_of_edges())
    print ""

    # make modules
    net, cliques = make_modules(net)
    print "Modules Formed"
    print "number of modules: " + str(len(cliques))
    print ""

    # print network
    nx.write_gml(net, 'conetwork.gml')

    # collapse modules
    coll_table = collapse_modules(table, cliques, args.prefix)
    print "Table Collapsed"

    # print new table
    coll_table.to_json('make_modules.py', open('collapsed.biom', 'w'))

if __name__ == '__main__':
    """main, takes argparser"""
    import argparse

    parser = argparse.ArgumentParser()
    parser.add_argument("-i", "--input", help="location of input biom file")
    parser.add_argument("-o", "--output", help="output file location")
    parser.add_argument("-m", "--correl_method", help="correlation method", default="spearman")
    parser.add_argument("-a", "--p_adjust", help="p-value adjustment", default="bh")
    parser.add_argument("-s", "--min_sample", help="minimum number of samples present in", type=int)
    parser.add_argument("--prefix", help="prefix for module names in collapsed file", default="module_")
    parser.add_argument("-k", "--k_size", help="desired k-size to determine cliques", default=3, type=int)
    parser.add_argument("--min_p", help="minimum p-value to determine edges", default=.05, type=float)
    parser.add_argument("--outlier_removal", help="outlier detection and removal", default=False, action="store_false")
    module_maker(parser.parse_args())
<commit_msg>removing tests from collapse modules, still don't work<commit_after>"""Make modules of observations based on cooccurence networks and collapse table"""

import os
import shutil
import glob
import sys

import networkx as nx
import numpy as np
import pysurvey as ps

import general
from biom import load_table, Table
from scipy.stats import spearmanr, pearsonr
from pysurvey import SparCC as sparcc
from functools import partial

def paired_correlations_from_table(table, correl_method=spearmanr, p_adjust=general.bh_adjust):
    """Takes a biom table and finds correlations between all pairs of otus."""
    correls = list()

    for (data_i, otu_i, metadata_i), (data_j, otu_j, metadata_j) in table.iter_pairwise(axis='observation'):
        correl = correl_method(data_i, data_j)
        correls.append([str(otu_i), str(otu_j), correl[0], correl[1]])

    # adjust p-value if desired
    if p_adjust is not None:
        p_adjusted = p_adjust([i[3] for i in correls])
        for i in xrange(len(correls)):
            correls[i].append(p_adjusted[i])

    header = ['feature1', 'feature2', 'r', 'p']
    if p_adjust is not None:
        header.append('adjusted_p')

    return correls, header

def boostrapped_correlation_multi(tup, temp_folder, cor_temp):
    """tup is a tuple where tup[0] is num and tup[1] is the file name"""
    df = ps.read_txt(tup[1], verbose=False)
    cor = ps.basis_corr(df, oprint=False)[0]
    ps.write_txt(cor, temp_folder+cor_temp+str(tup[0])+".txt", T=False)

def sparcc_correlations_multi(table, p_adjust=general.bh_adjust, temp_folder=os.getcwd()+"/temp/",
                              boot_temp="bootstrap_", cor_temp="cor_", table_temp="temp_table.txt",
                              bootstraps=100):
    """Calculate correlations with sparcc"""
    import multiprocessing

    os.mkdir(temp_folder)
    pool = multiprocessing.Pool(multiprocessing.cpu_count()-1)

    # make tab delimited, delete first line and reread in
    with open(temp_folder+table_temp, 'w') as f:
        f.write('\n'.join(table.to_tsv().split("\n")[1:]))
    df = ps.read_txt(temp_folder+table_temp, verbose=False)

    sparcc.make_bootstraps(df, bootstraps, boot_temp+"#.txt", temp_folder)
    cor, cov = ps.basis_corr(df, oprint=False)

    pfun = partial(boostrapped_correlation_multi, temp_folder=temp_folder, cor_temp=cor_temp)
    tups = enumerate(glob.glob(temp_folder+boot_temp+"*.txt"))
    pool.map(pfun, tups)
    pool.close()
    pool.join()

    p_vals = sparcc.get_pvalues(cor, temp_folder+cor_temp+"#.txt", bootstraps)
    # generate correls
    correls = list()
    for i in xrange(len(cor.index)):
        for j in xrange(i+1, len(cor.index)):
            correls.append([str(cor.index[i]), str(cor.index[j]), cor.iat[i, j], p_vals.iat[i, j]])

    # adjust p-value if desired
    if p_adjust is not None:
        p_adjusted = p_adjust([i[3] for i in correls])
        for i in xrange(len(correls)):
            correls[i].append(p_adjusted[i])

    header = ['feature1', 'feature2', 'r', 'p']
    if p_adjust is not None:
        header.append('adjusted_p')

    # cleanup, remove all of bootstraps folder
    shutil.rmtree(temp_folder)

    return correls, header

def boostrapped_correlation(in_file, temp_folder, cor_temp, num):
    df = ps.read_txt(in_file, verbose=False)
    cor = ps.basis_corr(df, oprint=False)[0]
    ps.write_txt(cor, temp_folder+cor_temp+str(num)+".txt", T=False)

def sparcc_correlations(table, p_adjust=general.bh_adjust):
    """"""
    # setup
    temp_folder = "temp3/"
    boot_temp = "bootstrap_"
    cor_temp = "cor_"
    table_temp = "temp_table.txt"
    os.mkdir(temp_folder)
    bootstraps = 100

    # make tab delimited, delete first line and reread in
    with open(temp_folder+table_temp, 'w') as f:
        f.write('\n'.join(table.to_tsv().split("\n")[1:]))
    df = ps.read_txt(temp_folder+table_temp, verbose=False)

    # calculate correlations
    cor, cov = ps.basis_corr(df, oprint=False)

    # calculate p-values
    sparcc.make_bootstraps(df, bootstraps, boot_temp+"#.txt", temp_folder)
    for i, file in enumerate(glob.glob(temp_folder+boot_temp+"*.txt")):
        boostrapped_correlation(file, temp_folder, cor_temp, i)

    p_vals = sparcc.get_pvalues(cor, temp_folder+cor_temp+"#.txt", bootstraps)
    # generate correls
    correls = list()
    for i in xrange(len(cor.index)):
        for j in xrange(i+1, len(cor.index)):
            correls.append([str(cor.index[i]), str(cor.index[j]), cor.iat[i, j], p_vals.iat[i, j]])

    # adjust p-value if desired
    if p_adjust is not None:
        p_adjusted = p_adjust([i[3] for i in correls])
        for i in xrange(len(correls)):
            correls[i].append(p_adjusted[i])

    header = ['feature1', 'feature2', 'r', 'p']
    if p_adjust is not None:
        header.append('adjusted_p')

    # cleanup, remove all of bootstraps folder
    shutil.rmtree(temp_folder)

    return correls, header

def paired_correlations_from_table_with_outlier_removal(table, good_samples, min_keep=10, correl_method=spearmanr,
                                                        p_adjust=general.bh_adjust):
    """Takes a biom table and finds correlations between all pairs of otus."""
    correls = list()

    for (data_i, otu_i, metadata_i), (data_j, otu_j, metadata_j) in table.iter_pairwise(axis='observation'):
        samp_union = np.union1d(good_samples[otu_i], good_samples[otu_j])
        # remove zero zero points
        # samp_union = [ind for i, ind in enumerate(samp_union) if data_i[i]!=0 and data_j[i]!=0]
        if len(samp_union) > min_keep:
            correl = correl_method(data_i[samp_union], data_j[samp_union])
            correls.append([str(otu_i), str(otu_j), correl[0], correl[1]])

    # adjust p-value if desired
    if p_adjust is not None:
        p_adjusted = p_adjust([i[3] for i in correls])
        for i in xrange(len(correls)):
            correls[i].append(p_adjusted[i])

    header = ['feature1', 'feature2', 'r', 'p']
    if p_adjust is not None:
        header.append('adjusted_p')

    return correls, header

def make_modules(graph, k=3):
    """make modules with networkx k-clique communities and annotate network
    :type graph: networkx graph
    """
    cliques = list(nx.k_clique_communities(graph, k))
    cliques = [list(i) for i in cliques]
    for i, clique in enumerate(cliques):
        for node in clique:
            graph.node[node]['k_' + str(k)] = i
    return graph, cliques

def make_modules_multik(graph, k=None):
    """make modules with networkx k-clique communities and annotate network"""
    if k is None:
        k = [2, 3, 4, 5, 6]
    communities = dict()
    for k_val in list(k):
        cliques = list(nx.k_clique_communities(graph, k_val))
        cliques = [list(i) for i in cliques]
        communities[k_val] = cliques
        for i, clique in enumerate(cliques):
            for node in clique:
                graph.node[node]['k_' + str(k_val)] = i
    return graph, {k: list(v) for k, v in communities.iteritems()}

def collapse_modules(table, cliques, prefix="module_"):
    """collapse created modules in a biom table"""
    in_module = set()
    modules = np.zeros((len(cliques), table.shape[1]))

    # for each clique merge values and print cliques to file
    with open("cliques.txt", 'w') as f:
        for i, clique in enumerate(cliques):
            in_module = in_module | set(clique)
            f.write(prefix+str(i)+'\t'+','.join([str(i) for i in clique])+'\n')
            for feature in clique:
                try:
                    modules[i] += table.data(feature, axis="observation")
                except:
                    print feature
                    print feature in table.ids(axis="observation")
                    sys.exit()
    table.filter(in_module, axis='observation')

    # make new table
    new_table_matrix = np.concatenate((table.matrix_data.toarray(), modules))
    new_table_obs = list(table.ids(axis='observation')) + [prefix + str(i) for i in range(0, len(cliques))]
    return Table(new_table_matrix, new_table_obs, table.ids())

def collapse_modules_multik(table, cliques, prefix="module_"):
    """collapse created modules in a biom table"""
    in_module = set()
    modules = np.zeros((len(cliques), table.shape[1]))

    # for each clique merge values
    for i, clique in enumerate(cliques):
        in_module = in_module | set(clique)
        for feature in clique:
            modules[i] += table.data(feature, axis='observation')
    table.filter(in_module, axis='observation')

    # make new table
    new_table_matrix = np.concatenate((table.matrix_data.toarray(), modules))
    new_table_obs = list(table.ids(axis='observation')) + [prefix + str(i) for i in range(0, len(cliques))]
    return Table(new_table_matrix, new_table_obs, table.ids())

def module_maker(args):
    # correlation and p-value adjustment methods
    correl_methods = {'spearman': spearmanr, 'pearson': pearsonr, 'sparcc': sparcc_correlations_multi}
    p_methods = {'bh': general.bh_adjust, 'bon': general.bonferroni_adjust}
    correl_method = correl_methods[args.correl_method]
    if args.p_adjust is not None:
        p_adjust = p_methods[args.p_adjust]
    else:
        p_adjust = None

    # get features to be correlated and extract metadata
    table = load_table(args.input)
    metadata = general.get_metadata_from_table(table)
    print "Table loaded: " + str(table.shape[0]) + " observations"
    print ""

    # make new output directory and change to it
    if args.output is not None:
        os.makedirs(args.output)
        os.chdir(args.output)

    # convert to relative abundance and filter
    if args.min_sample != None:
        table_filt = general.filter_table(table, args.min_sample)
        print "Table filtered: " + str(table_filt.shape[0]) + " observations"
        print ""
    else:
        table_filt = table

    # correlate feature
    if correl_method in [spearmanr, pearsonr]:
        if args.outlier_removal:
            # remove outlier observations
            # first attempt with just looking at individual otu's
            good_samples = general.remove_outliers(table_filt)
            print "Outliers removed: " + str(len(good_samples)) + " observations"
            print ""
            correls, correl_header = paired_correlations_from_table_with_outlier_removal(table_filt, good_samples,
                                                                                correl_method, p_adjust)
        else:
            # correlate feature
            correls, correl_header = paired_correlations_from_table(table_filt, correl_method, p_adjust)
    else:
        correls, correl_header = sparcc_correlations(table_filt, p_adjust)
    general.print_delimited('correls.txt', correls, correl_header)

    print "Features Correlated"

    # make correlation network
    net = general.correls_to_net(correls, conet=True, metadata=metadata, min_p=args.min_p)
    print "Network Generated"
    print "number of nodes: " + str(net.number_of_nodes())
    print "number of edges: " + str(net.number_of_edges())
    print ""

    # make modules
    net, cliques = make_modules(net)
    print "Modules Formed"
    print "number of modules: " + str(len(cliques))
    print ""

    # print network
    nx.write_gml(net, 'conetwork.gml')

    # collapse modules
    coll_table = collapse_modules(table, cliques, args.prefix)
    print "Table Collapsed"

    # print new table
    coll_table.to_json('make_modules.py', open('collapsed.biom', 'w'))

if __name__ == '__main__':
    """main, takes argparser"""
    import argparse

    parser = argparse.ArgumentParser()
    parser.add_argument("-i", "--input", help="location of input biom file")
    parser.add_argument("-o", "--output", help="output file location")
    parser.add_argument("-m", "--correl_method", help="correlation method", default="spearman")
    parser.add_argument("-a", "--p_adjust", help="p-value adjustment", default="bh")
    parser.add_argument("-s", "--min_sample", help="minimum number of samples present in", type=int)
    parser.add_argument("--prefix", help="prefix for module names in collapsed file", default="module_")
    parser.add_argument("-k", "--k_size", help="desired k-size to determine cliques", default=3, type=int)
    parser.add_argument("--min_p", help="minimum p-value to determine edges", default=.05, type=float)
    parser.add_argument("--outlier_removal", help="outlier detection and removal", default=False, action="store_false")
    module_maker(parser.parse_args())
<|endoftext|>