<commit_before>import unittest

from io import BytesIO, BufferedReader

from urllib3.response import HTTPResponse
from urllib3.exceptions import DecodeError

class TestLegacyResponse(unittest.TestCase):
    def test_getheaders(self):
        headers = {'host': 'example.com'}
        r = HTTPResponse(headers=headers)
        self.assertEqual(r.getheaders(), headers)

    def test_getheader(self):
        headers = {'host': 'example.com'}
        r = HTTPResponse(headers=headers)
        self.assertEqual(r.getheader('host'), 'example.com')

class TestResponse(unittest.TestCase):
    def test_cache_content(self):
        r = HTTPResponse('foo')
        self.assertEqual(r.data, 'foo')
        self.assertEqual(r._body, 'foo')

    def test_default(self):
        r = HTTPResponse()
        self.assertEqual(r.data, None)

    def test_none(self):
        r = HTTPResponse(None)
        self.assertEqual(r.data, None)

    def test_preload(self):
        fp = BytesIO(b'foo')

        r = HTTPResponse(fp, preload_content=True)

        self.assertEqual(fp.tell(), len(b'foo'))
        self.assertEqual(r.data, b'foo')

    def test_no_preload(self):
        fp = BytesIO(b'foo')

        r = HTTPResponse(fp, preload_content=False)

        self.assertEqual(fp.tell(), 0)
        self.assertEqual(r.data, b'foo')
        self.assertEqual(fp.tell(), len(b'foo'))

    def test_decode_bad_data(self):
        fp = BytesIO(b'\x00' * 10)
        self.assertRaises(DecodeError, HTTPResponse, fp, headers={
            'content-encoding': 'deflate'
        })

    def test_decode_deflate(self):
        import zlib
        data = zlib.compress(b'foo')

        fp = BytesIO(data)
        r = HTTPResponse(fp, headers={'content-encoding': 'deflate'})

        self.assertEqual(r.data, b'foo')

    def test_decode_deflate_case_insensitve(self):
        import zlib
        data = zlib.compress(b'foo')

        fp = BytesIO(data)
        r = HTTPResponse(fp, headers={'content-encoding': 'DeFlAtE'})

        self.assertEqual(r.data, b'foo')

    def test_chunked_decoding_deflate(self):
        import zlib
        data = zlib.compress(b'foo')

        fp = BytesIO(data)
        r = HTTPResponse(fp, headers={'content-encoding': 'deflate'},
                         preload_content=False)

        self.assertEqual(r.read(3), b'')
        self.assertEqual(r.read(1), b'f')
        self.assertEqual(r.read(2), b'oo')

    def test_chunked_decoding_deflate2(self):
        import zlib
        compress = zlib.compressobj(6, zlib.DEFLATED, -zlib.MAX_WBITS)
        data = compress.compress(b'foo')
        data += compress.flush()

        fp = BytesIO(data)
        r = HTTPResponse(fp, headers={'content-encoding': 'deflate'},
                         preload_content=False)

        self.assertEqual(r.read(1), b'')
        self.assertEqual(r.read(1), b'f')
        self.assertEqual(r.read(2), b'oo')

    def test_chunked_decoding_gzip(self):
        import zlib
        compress = zlib.compressobj(6, zlib.DEFLATED, 16 + zlib.MAX_WBITS)
        data = compress.compress(b'foo')
        data += compress.flush()

        fp = BytesIO(data)
        r = HTTPResponse(fp, headers={'content-encoding': 'gzip'},
                         preload_content=False)

        self.assertEqual(r.read(11), b'')
        self.assertEqual(r.read(1), b'f')
        self.assertEqual(r.read(2), b'oo')

    def test_io(self):
        import socket
        try:
            from http.client import HTTPResponse as OldHTTPResponse
        except:
            from httplib import HTTPResponse as OldHTTPResponse

        fp = BytesIO(b'foo')
        resp = HTTPResponse(fp, preload_content=False)

        self.assertEqual(resp.closed, False)
        self.assertEqual(resp.readable(), True)
        self.assertEqual(resp.writable(), False)
        self.assertRaises(IOError, resp.fileno)

        resp.close()
        self.assertEqual(resp.closed, True)

        # Try closing with an `httplib.HTTPResponse`, because it has an
        # `isclosed` method.
        hlr = OldHTTPResponse(socket.socket())
        resp2 = HTTPResponse(hlr, preload_content=False)
        self.assertEqual(resp2.closed, False)
        resp2.close()
        self.assertEqual(resp2.closed, True)

        #also try when only data is present.
        resp3 = HTTPResponse('foodata')
        self.assertRaises(IOError, resp3.fileno)

        resp3._fp = 2
        # A corner case where _fp is present but doesn't have `closed`,
        # `isclosed`, or `fileno`.  Unlikely, but possible.
        self.assertEqual(resp3.closed, True)
        self.assertRaises(IOError, resp3.fileno)

    def test_io_bufferedreader(self):
        fp = BytesIO(b'foo')
        resp = HTTPResponse(fp, preload_content=False)
        br = BufferedReader(resp)

        self.assertEqual(br.read(), b'foo')

        br.close()
        self.assertEqual(resp.closed, True)

    def test_streaming(self):
        fp = BytesIO(b'foo')
        resp = HTTPResponse(fp, preload_content=False)
        stream = resp.stream(2, decode_content=False)

        self.assertEqual(next(stream), b'fo')
        self.assertEqual(next(stream), b'o')
        self.assertRaises(StopIteration, next, stream)

    def test_gzipped_streaming(self):
        import zlib
        compress = zlib.compressobj(6, zlib.DEFLATED, 16 + zlib.MAX_WBITS)
        data = compress.compress(b'foo')
        data += compress.flush()

        fp = BytesIO(data)
        resp = HTTPResponse(fp, headers={'content-encoding': 'gzip'},
                         preload_content=False)
        stream = resp.stream(2)

        self.assertEqual(next(stream), b'f')
        self.assertEqual(next(stream), b'oo')
        self.assertRaises(StopIteration, next, stream)

    def test_deflate_streaming(self):
        import zlib
        data = zlib.compress(b'foo')

        fp = BytesIO(data)
        resp = HTTPResponse(fp, headers={'content-encoding': 'deflate'},
                         preload_content=False)
        stream = resp.stream(2)

        self.assertEqual(next(stream), b'f')
        self.assertEqual(next(stream), b'oo')
        self.assertRaises(StopIteration, next, stream)

    def test_deflate2_streaming(self):
        import zlib
        compress = zlib.compressobj(6, zlib.DEFLATED, -zlib.MAX_WBITS)
        data = compress.compress(b'foo')
        data += compress.flush()

        fp = BytesIO(data)
        resp = HTTPResponse(fp, headers={'content-encoding': 'deflate'},
                         preload_content=False)
        stream = resp.stream(2)

        self.assertEqual(next(stream), b'f')
        self.assertEqual(next(stream), b'oo')
        self.assertRaises(StopIteration, next, stream)

    def test_empty_stream(self):
        fp = BytesIO(b'')
        resp = HTTPResponse(fp, preload_content=False)
        stream = resp.stream(2, decode_content=False)

        self.assertRaises(StopIteration, next, stream)

if __name__ == '__main__':
    unittest.main()
<commit_msg>Test with a mocked HTTPResponse object.<commit_after>import unittest

from io import BytesIO, BufferedReader

from urllib3.response import HTTPResponse
from urllib3.exceptions import DecodeError

class TestLegacyResponse(unittest.TestCase):
    def test_getheaders(self):
        headers = {'host': 'example.com'}
        r = HTTPResponse(headers=headers)
        self.assertEqual(r.getheaders(), headers)

    def test_getheader(self):
        headers = {'host': 'example.com'}
        r = HTTPResponse(headers=headers)
        self.assertEqual(r.getheader('host'), 'example.com')

class TestResponse(unittest.TestCase):
    def test_cache_content(self):
        r = HTTPResponse('foo')
        self.assertEqual(r.data, 'foo')
        self.assertEqual(r._body, 'foo')

    def test_default(self):
        r = HTTPResponse()
        self.assertEqual(r.data, None)

    def test_none(self):
        r = HTTPResponse(None)
        self.assertEqual(r.data, None)

    def test_preload(self):
        fp = BytesIO(b'foo')

        r = HTTPResponse(fp, preload_content=True)

        self.assertEqual(fp.tell(), len(b'foo'))
        self.assertEqual(r.data, b'foo')

    def test_no_preload(self):
        fp = BytesIO(b'foo')

        r = HTTPResponse(fp, preload_content=False)

        self.assertEqual(fp.tell(), 0)
        self.assertEqual(r.data, b'foo')
        self.assertEqual(fp.tell(), len(b'foo'))

    def test_decode_bad_data(self):
        fp = BytesIO(b'\x00' * 10)
        self.assertRaises(DecodeError, HTTPResponse, fp, headers={
            'content-encoding': 'deflate'
        })

    def test_decode_deflate(self):
        import zlib
        data = zlib.compress(b'foo')

        fp = BytesIO(data)
        r = HTTPResponse(fp, headers={'content-encoding': 'deflate'})

        self.assertEqual(r.data, b'foo')

    def test_decode_deflate_case_insensitve(self):
        import zlib
        data = zlib.compress(b'foo')

        fp = BytesIO(data)
        r = HTTPResponse(fp, headers={'content-encoding': 'DeFlAtE'})

        self.assertEqual(r.data, b'foo')

    def test_chunked_decoding_deflate(self):
        import zlib
        data = zlib.compress(b'foo')

        fp = BytesIO(data)
        r = HTTPResponse(fp, headers={'content-encoding': 'deflate'},
                         preload_content=False)

        self.assertEqual(r.read(3), b'')
        self.assertEqual(r.read(1), b'f')
        self.assertEqual(r.read(2), b'oo')

    def test_chunked_decoding_deflate2(self):
        import zlib
        compress = zlib.compressobj(6, zlib.DEFLATED, -zlib.MAX_WBITS)
        data = compress.compress(b'foo')
        data += compress.flush()

        fp = BytesIO(data)
        r = HTTPResponse(fp, headers={'content-encoding': 'deflate'},
                         preload_content=False)

        self.assertEqual(r.read(1), b'')
        self.assertEqual(r.read(1), b'f')
        self.assertEqual(r.read(2), b'oo')

    def test_chunked_decoding_gzip(self):
        import zlib
        compress = zlib.compressobj(6, zlib.DEFLATED, 16 + zlib.MAX_WBITS)
        data = compress.compress(b'foo')
        data += compress.flush()

        fp = BytesIO(data)
        r = HTTPResponse(fp, headers={'content-encoding': 'gzip'},
                         preload_content=False)

        self.assertEqual(r.read(11), b'')
        self.assertEqual(r.read(1), b'f')
        self.assertEqual(r.read(2), b'oo')

    def test_io(self):
        import socket
        try:
            from http.client import HTTPResponse as OldHTTPResponse
        except:
            from httplib import HTTPResponse as OldHTTPResponse

        fp = BytesIO(b'foo')
        resp = HTTPResponse(fp, preload_content=False)

        self.assertEqual(resp.closed, False)
        self.assertEqual(resp.readable(), True)
        self.assertEqual(resp.writable(), False)
        self.assertRaises(IOError, resp.fileno)

        resp.close()
        self.assertEqual(resp.closed, True)

        # Try closing with an `httplib.HTTPResponse`, because it has an
        # `isclosed` method.
        hlr = OldHTTPResponse(socket.socket())
        resp2 = HTTPResponse(hlr, preload_content=False)
        self.assertEqual(resp2.closed, False)
        resp2.close()
        self.assertEqual(resp2.closed, True)

        #also try when only data is present.
        resp3 = HTTPResponse('foodata')
        self.assertRaises(IOError, resp3.fileno)

        resp3._fp = 2
        # A corner case where _fp is present but doesn't have `closed`,
        # `isclosed`, or `fileno`.  Unlikely, but possible.
        self.assertEqual(resp3.closed, True)
        self.assertRaises(IOError, resp3.fileno)

    def test_io_bufferedreader(self):
        fp = BytesIO(b'foo')
        resp = HTTPResponse(fp, preload_content=False)
        br = BufferedReader(resp)

        self.assertEqual(br.read(), b'foo')

        br.close()
        self.assertEqual(resp.closed, True)

    def test_streaming(self):
        fp = BytesIO(b'foo')
        resp = HTTPResponse(fp, preload_content=False)
        stream = resp.stream(2, decode_content=False)

        self.assertEqual(next(stream), b'fo')
        self.assertEqual(next(stream), b'o')
        self.assertRaises(StopIteration, next, stream)

    def test_gzipped_streaming(self):
        import zlib
        compress = zlib.compressobj(6, zlib.DEFLATED, 16 + zlib.MAX_WBITS)
        data = compress.compress(b'foo')
        data += compress.flush()

        fp = BytesIO(data)
        resp = HTTPResponse(fp, headers={'content-encoding': 'gzip'},
                         preload_content=False)
        stream = resp.stream(2)

        self.assertEqual(next(stream), b'f')
        self.assertEqual(next(stream), b'oo')
        self.assertRaises(StopIteration, next, stream)

    def test_deflate_streaming(self):
        import zlib
        data = zlib.compress(b'foo')

        fp = BytesIO(data)
        resp = HTTPResponse(fp, headers={'content-encoding': 'deflate'},
                         preload_content=False)
        stream = resp.stream(2)

        self.assertEqual(next(stream), b'f')
        self.assertEqual(next(stream), b'oo')
        self.assertRaises(StopIteration, next, stream)

    def test_deflate2_streaming(self):
        import zlib
        compress = zlib.compressobj(6, zlib.DEFLATED, -zlib.MAX_WBITS)
        data = compress.compress(b'foo')
        data += compress.flush()

        fp = BytesIO(data)
        resp = HTTPResponse(fp, headers={'content-encoding': 'deflate'},
                         preload_content=False)
        stream = resp.stream(2)

        self.assertEqual(next(stream), b'f')
        self.assertEqual(next(stream), b'oo')
        self.assertRaises(StopIteration, next, stream)

    def test_empty_stream(self):
        fp = BytesIO(b'')
        resp = HTTPResponse(fp, preload_content=False)
        stream = resp.stream(2, decode_content=False)

        self.assertRaises(StopIteration, next, stream)

    def test_mock_httpresponse_stream(self):
        # Mock out a HTTP Request that does enough to make it through urllib3's
        # read() and close() calls, and also exhausts and underlying file
        # object.
        class MockHTTPRequest(object):
            self.fp = None

            def read(self, amt):
                data = self.fp.read(amt)
                if not data:
                    self.fp = None

                return data

            def close(self):
                self.fp = None

        bio = BytesIO(b'foo')
        fp = MockHTTPRequest()
        fp.fp = bio
        resp = HTTPResponse(fp, preload_content=False)
        stream = resp.stream(2)

        self.assertEqual(next(stream), b'fo')
        self.assertEqual(next(stream), b'o')
        self.assertRaises(StopIteration, next, stream)

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


<commit_before># -*- coding: utf-8 -*-
"""
Created on Sun Apr 16 14:32:23 2017

@author: scram
"""
from bson.son import SON
from pymongo import MongoClient
import networkx as nx
import bson

test_dict ={
851588099360649216: "root_tweet",
851588741055623169: "top_reply_1",
851613881520803840: "reply_1_1",
851757634613018625: "reply_1_2",
851759076577079296: "reply_1_2_1",
851595539565105152: "top_reply_2",
851596235949658112: "reply_2_1",
851782925981028353: "reply_2_1_1",
851935916725764097: "reply_2_1_2",
851595212468084736: "top_reply_3"
 }

def get_tweet_time(TWEET):
    strfts = time.strftime('%Y-%m-%d %H:%M:%S', time.strptime(TWEET['created_at'],'%a %b %d %H:%M:%S +0000 %Y'))
    ts = time.strftime('%Y-%m-%d %H:%M:%S',time.gmtime(time.mktime(time.strptime(TWEET['created_at'],'%a %b %d %H:%M:%S +0000 %Y'))) )
    return ts,strfts

def id2edgelist(mongo_id,collection):
    list_dicts = list(collection.find({"_id":biggest_id}))[0]
    edges = []
    for i in list_dicts['edge_list']:
        edges.append((i['parent'],i['child']))
    return edges

def edge_list(collections_list,DB):
    """ Takes Tweets from MongoDB and creates directed edge_lists in new collections

    ARGS(list: str) Mongo collection names in the order [children, parents, ...]
    ARGS(object: pymongo.database.Database) Mongo database with the collecctions

    """
    #loop though the collections
    child_coll = collections_list[0]
    parent_coll = collections_list[1]

    edge_coll = db.edge_list
    root_ids = []
    child_ids = set(edge_coll.distinct("edge_list.child"))
    parent_ids = set(edge_coll.distinct("edge_list.parent"))
    root_ids = set(edge_coll.distinct("_id"))

    print root_ids == parent_ids
    for num, collection in enumerate(collections_list):
        if num == 0:
            print collection
            for tweet in DB[collection].find({}).sort("_id",1):
                edge = [tweet["in_reply_to_status_id"],tweet["id"]]
                if edge[1] not in child_ids:
                    if edge[0]:
                        if edge[0] in child_ids:
                            edge_coll.update_one(
                                    {"edge_list.child":edge[0]},
                                    {"$push": {"edge_list":
                                        {"parent":edge[0],
                                         "child":edge[1]}
                                        }
                                    })
                            child_ids.add(edge[1])
                            parent_ids.add(edge[0])
                            at_root = list(edge_coll.find({"edge_list.child":edge[0]}))[0]["_id"]
                            print at_root,"^--",edge[0],"<--",edge[1]
    #                        if edge[1] in root_ids:
    #                            subThread =list(edge_coll.find({"_id":edge[1]}))
    #                            for edg in subThread[0]["edge_list"]:
    #                                edge_coll.update_one(
    #                                    {"edge_list.child":edge[0]},
    #                                    {"$push": {"edge_list":
    #                                        edg
    #                                        }
    #                                    })
    #                            edge_coll.delete_one({"_id":edge[1]})
                        elif edge[0] not in root_ids:
                            if db[parent_coll].find({'_id':edge[0]}):
                                post = {"_id": edge[0],
                                    "edge_list": [
                                        {"parent":edge[0],
                                        "child":edge[1]}
                                        ]}
                                post_id = edge_coll.insert_one(post).inserted_id
                                print "^",post_id,"<--",edge[1]
                                root_ids.add(post_id)
                                child_ids.add(edge[1])
                                parent_ids.add(edge[0])
                            else:
                                db[child_coll].delete_one({'_id':tweet['id']})

                        elif edge[1] not in child_ids:
                            edge_coll.update_one(
                                    {"_id":edge[0]},
                                    {"$push": {"edge_list":
                                        {"parent":edge[0],
                                         "child":edge[1]}
                                        }
                                    })
                            child_ids.add(edge[1])
                            parent_ids.add(edge[0])
                            print edge[0],"^--",edge[1]
            print "Finished first collection"
        else:
            # Decide if you want to do anything to with this to support different
            # strage approches....
            threads = 0
            print collection
            print "\t",collections_list[num-1]
            children = collections_list[num-1]
            child_ids = DB[children].distinct("id")
            for tweet in DB[collection].find():
                if edge[0]:
                    if edge[1] in child_ids:
                        threads += 1
            print "threads", threads

def remove_old_threads(DB,id_date):
    test =856172056932700164
#    for thread in list(DB.edge_list.find({"_id":test}))[:10]:
    del_replies = 0
    del_roots = 0
    del_threads = 0
    for thread in list(DB.edge_list.find({}))[:100]:
        thread_old = []
        id_set = set()
        if thread['_id'] < id_date:
            for edge in thread['edge_list']:
                id_set.add(edge['parent'])
                id_set.add(edge['child'])
            del_roots +=1
            id_list = list(id_set - set([thread['_id']]))
            del_replies += len(id_list)
            DB.trump_tweets.delete_one({'_id':thread['_id']})
            del_dic ={'_id':{'$in':list(id_set)}}
            DB.replies_to_trump.delete_many(del_dic)
            DB.edge_list.delete_one({'_id':thread['_id']})
            del_threads +=1
        else:
            for edge in thread['edge_list']:
                id_set.add(edge['parent'])
                id_set.add(edge['child'])
                thread_old.append(edge['child'] < id_date)
                thread_old.append(edge['parent'] < id_date)
            if any(thread_old):
                del_roots +=1
                id_list = list(id_set - set([thread['_id']]))
                del_replies += len(id_list)
                DB.trump_tweets.delete_one({'_id':thread['_id']})
                del_dic ={'_id':{'$in':list(id_set)}}
                DB.replies_to_trump.delete_many(del_dic)
                DB.edge_list.delete_one({'_id':thread['_id']})
                del_threads +=1

    changes = [del_replies,del_roots,del_threads]
    if sum(changes)>0:
        print 'Number deleted from "replies_to_tweets"',del_replies
        print 'Number deleted from "trump_tweets"',del_roots
        print 'Number deleted from "edge_list"',del_threads
    else:
        print "None deleted"
    print "FINISHED"

if __name__ == "__main__":
    client = MongoClient()
    client = MongoClient('localhost', 27017)
#    db = client.Alta_Real
#    db = client["test-tree"]
    db = client.test_tree
    collection_names = ["replies_to_trump","trump_tweets","test_parents"]
    collection_names = ["replies_to_trump","trump_tweets"]

#    """
    edge_list(collection_names,db)
    edge_collection = db.edge_list
    big_thread = edge_collection.aggregate( [{ "$unwind" : "$edge_list" },
            { "$group" : { "_id" : "$_id", "len" : { "$sum" : 1 } } },
            { "$sort" : { "len" : -1 } },
            { "$limit" : 1 }
            ] )
#    print list(big_thread)
    biggest_id =  list(big_thread)[0]["_id"]
    print biggest_id
    DG=nx.DiGraph()
    DG.add_edges_from(id2edgelist(biggest_id,edge_collection))
    nx.draw_random(DG, with_labels=False)
    nx.draw_random(DG, with_labels=True)
#    """
    old_id = 856425255862235136
#    oldest_time = list(db.replies_to_trump.find({"_id":old_id}))[0]["created_at"]
#    oldest_id = list(db.replies_to_trump.find({"_id":old_id}))[0]["_id"]
    remove_old_threads(db,old_id)
#    remove_old_threads(db,oldest_time)

    <commit_msg>FINAL COMMIT bug fixes that I can't remeber<commit_after># -*- coding: utf-8 -*-
"""
Created on Sun Apr 16 14:32:23 2017

@author: scram
"""
from bson.son import SON
from pymongo import MongoClient
import networkx as nx
import bson

test_dict ={
851588099360649216: "root_tweet",
851588741055623169: "top_reply_1",
851613881520803840: "reply_1_1",
851757634613018625: "reply_1_2",
851759076577079296: "reply_1_2_1",
851595539565105152: "top_reply_2",
851596235949658112: "reply_2_1",
851782925981028353: "reply_2_1_1",
851935916725764097: "reply_2_1_2",
851595212468084736: "top_reply_3"
 }

def get_tweet_time(TWEET):
    strfts = time.strftime('%Y-%m-%d %H:%M:%S', time.strptime(TWEET['created_at'],'%a %b %d %H:%M:%S +0000 %Y'))
    ts = time.strftime('%Y-%m-%d %H:%M:%S',time.gmtime(time.mktime(time.strptime(TWEET['created_at'],'%a %b %d %H:%M:%S +0000 %Y'))) )
    return ts,strfts

def id2edgelist(mongo_id,collection):
    list_dicts = list(collection.find({"_id":mongo_id}))[0]
    edges = []
    for i in list_dicts['edge_list']:
        edges.append((i['parent'],i['child']))
    return edges

def edge_list(collections_list,DB):
    """ Takes Tweets from MongoDB and creates directed edge_lists in new collections

    ARGS(list: str) Mongo collection names in the order [children, parents, ...]
    ARGS(object: pymongo.database.Database) Mongo database with the collecctions

    """
    #loop though the collections
    child_coll = collections_list[0]
    parent_coll = collections_list[1]

    edge_coll = db.edge_list
    root_ids = []
    child_ids = set(edge_coll.distinct("edge_list.child"))
    parent_ids = set(edge_coll.distinct("edge_list.parent"))
    root_ids = set(edge_coll.distinct("_id"))

    print(root_ids == parent_ids)
    for num, collection in enumerate(collections_list):
        if num == 0:
            print(collection)
            for tweet in DB[collection].find({}).sort("_id",1):
                edge = [tweet["in_reply_to_status_id"],tweet["id"]]
                if edge[1] not in child_ids:
                    if edge[0]:
                        if edge[0] in child_ids:
                            edge_coll.update_one(
                                    {"edge_list.child":edge[0]},
                                    {"$push": {"edge_list":
                                        {"parent":edge[0],
                                         "child":edge[1]}
                                        }
                                    })
                            child_ids.add(edge[1])
                            parent_ids.add(edge[0])
                            at_root = list(edge_coll.find({"edge_list.child":edge[0]}))[0]["_id"]
                            print( at_root,"^--",edge[0],"<--",edge[1])
                            if edge[1] in root_ids:
                                subThread =list(edge_coll.find({"_id":edge[1]}))
                                for edg in subThread[0]["edge_list"]:
                                    edge_coll.update_one(
                                        {"edge_list.child":edge[0]},
                                        {"$push": {"edge_list":
                                            edg
                                            }
                                        })
                                edge_coll.delete_one({"_id":edge[1]})
                        elif edge[0] not in root_ids:
                            if db[parent_coll].find({'_id':edge[0]}):
                                post = {"_id": edge[0],
                                    "edge_list": [
                                        {"parent":edge[0],
                                        "child":edge[1]}
                                        ]}
                                post_id = edge_coll.insert_one(post).inserted_id
                                print( "^",post_id,"<--",edge[1])
                                root_ids.add(post_id)
                                child_ids.add(edge[1])
                                parent_ids.add(edge[0])
                            else:
                                db[child_coll].delete_one({'_id':tweet['id']})

                        elif edge[1] not in child_ids:
                            edge_coll.update_one(
                                    {"_id":edge[0]},
                                    {"$push": {"edge_list":
                                        {"parent":edge[0],
                                         "child":edge[1]}
                                        }
                                    })
                            child_ids.add(edge[1])
                            parent_ids.add(edge[0])
                            print( edge[0],"^--",edge[1])
            print( "Finished first collection")
        else:
            # Decide if you want to do anything to with this to support different
            # strage approches....
            threads = 0
            print( collection)
            print( "\t",collections_list[num-1])
            children = (collections_list[num-1])
            child_ids = DB[children].distinct("id")
#            for tweet in DB[collection].find():
#                if edge[0]:
#                    if edge[1] in child_ids:
#                        threads += 1
            print( "threads", threads)

def remove_old_threads(DB,id_date):
    test =856172056932700164
#    for thread in list(DB.edge_list.find({"_id":test}))[:10]:
    del_replies = 0
    del_roots = 0
    del_threads = 0
    for thread in list(DB.edge_list.find({}))[:100]:
        thread_old = []
        id_set = set()
        if thread['_id'] < id_date:
            for edge in thread['edge_list']:
                id_set.add(edge['parent'])
                id_set.add(edge['child'])
            del_roots +=1
            id_list = list(id_set - set([thread['_id']]))
            del_replies += len(id_list)
            DB.trump_tweets.delete_one({'_id':thread['_id']})
            del_dic ={'_id':{'$in':list(id_set)}}
            DB.replies_to_trump.delete_many(del_dic)
            DB.edge_list.delete_one({'_id':thread['_id']})
            del_threads +=1
        else:
            for edge in thread['edge_list']:
                id_set.add(edge['parent'])
                id_set.add(edge['child'])
                thread_old.append(edge['child'] < id_date)
                thread_old.append(edge['parent'] < id_date)
            if any(thread_old):
                del_roots +=1
                id_list = list(id_set - set([thread['_id']]))
                del_replies += len(id_list)
                DB.trump_tweets.delete_one({'_id':thread['_id']})
                del_dic ={'_id':{'$in':list(id_set)}}
                DB.replies_to_trump.delete_many(del_dic)
                DB.edge_list.delete_one({'_id':thread['_id']})
                del_threads +=1

    changes = [del_replies,del_roots,del_threads]
    if sum(changes)>0:
        print( 'Number deleted from "replies_to_tweets"',del_replies)
        print( 'Number deleted from "trump_tweets"',del_roots)
        print( 'Number deleted from "edge_list"',del_threads)
    else:
        print( "None deleted")
    print("FINISHED")

if __name__ == "__main__":
    client = MongoClient()
    client = MongoClient('localhost', 27017)
#    db = client.Alta_Real
    db = client.Alta_Real_New

#    db = client["test-tree"]
#    db = client.test_tree
    collection_names = ["replies_to_trump","trump_tweets"]

#    """
    edge_list(collection_names,db)
    edge_collection = db.edge_list
    big_thread = list(edge_collection.aggregate( [{ "$unwind" : "$edge_list" },
            { "$group" : { "_id" : "$_id", "len" : { "$sum" : 1 } } },
            { "$sort" : { "len" : 1 } },
            { "$limit" : 1 }
            ] ))
#    print list(big_thread)
    print( len(big_thread) )
    biggest_id =  big_thread[0]["_id"]
    print( biggest_id)
    DG=nx.DiGraph()
    DG.add_edges_from(id2edgelist(biggest_id,edge_collection))
    nx.draw_random(DG, with_labels=False)
    nx.draw_random(DG, with_labels=True)
#    """
    pipeline_oldest =[{"$group":
                {"_id":"$in_reply_to_screen_name",
                "oldest":{"$min":"$id"},
                "newest":{"$max":"$id"},
                "count": {"$sum":1}
                }
            }]

#    oldest_time = list(db.replies_to_trump.find({"_id":old_id}))[0]["created_at"]
    oldest_id =list(db.replies_to_trump.aggregate(pipeline_oldest))[0]["oldest"]
    remove_old_threads(db,oldest_id)
#    remove_old_threads(db,oldest_time)

    <|endoftext|>


<commit_before>#!/usr/bin/env python2.7
# -*- coding: utf-8 -*-

import argparse
import sys
import os
import subprocess
import csv
import time
import datetime

import pymongo

# import colors

class VariantParseError(Exception):
    def __init__(self, value):
        self.value = value
    def __str__(self):
        return repr(self.value)

def import_variants(file_path, population, sex, file_format, user_id, mongo_port=27017):
    """doc"""

    # remove dots from file_name
    file_name = file_path.split('/')[-1]
    file_name_cleaned = file_name.replace('.', '').replace(' ', '')
    print >>sys.stderr, '[INFO] collection name: {}'.format(file_name_cleaned)

    with pymongo.Connection(port=mongo_port) as connection:
        db = connection['pergenie']
        users_variants = db['variants'][user_id][file_name_cleaned]
        data_info = db['data_info']

        data_info.update({'user_id': user_id, 'name': file_name_cleaned}, {"$set": {'file_name_cleaned': file_name_cleaned}}, upsert=True)

        # ensure this variants file is not imported
        if users_variants.find_one():
            db.drop_collection(users_variants)
            # print >>sys.stderr, colors.yellow('[WARNING] dropped old collection of {}'.format(file_name_cleaned))

        print >>sys.stderr, '[INFO] Countiong input lines ...',
        print >>sys.stderr, '[INFO] {}'.format(file_path)
        file_lines = int(subprocess.check_output(['wc', '-l', file_path]).split()[0])
        print >>sys.stderr, 'done. # of lines: {}'.format(file_lines)

        today = str(datetime.datetime.today()).replace('-', '/')
        info = {'user_id': user_id,
                'name': file_name_cleaned,
                'raw_name': file_name,
                'date': today,
                'population': population,
                'sex': sex,
                'file_format': file_format,
                'status': float(1.0)}

        data_info.update({'user_id': user_id, 'name': file_name_cleaned}, {"$set": info}, upsert=True)

        print >>sys.stderr,'[INFO] Start importing {0} as {1}...'.format(file_name, file_name_cleaned)
        prev_collections = db.collection_names()

        try:
            with open(file_path, 'rb') as fin:
                for i,data in enumerate(parse_lines(fin, file_format)):
                    if data['rs']:
                        users_variants.insert(data)

                    if i>0 and i%10000 == 0:
                        upload_status = int( 100 * ( i*0.9 / file_lines ) + 1 )
                        data_info.update({'user_id': user_id, 'name': file_name_cleaned}, {"$set": {'status': upload_status}})
                        print '[INFO] status: {}'.format(data_info.find({'user_id': user_id, 'name': file_name_cleaned})[0]['status'])

                print >>sys.stderr,'[INFO] ensure_index ...'
                users_variants.ensure_index('rs', unique=True)  # ok?

                data_info.update({'user_id': user_id, 'name': file_name_cleaned}, {"$set": {'status': 100}})

                print >>sys.stderr, '[INFO] status: {}'.format(data_info.find({'user_id': user_id, 'name': file_name_cleaned})[0]['status'])
                print >>sys.stderr, '[INFO] done. added collection {}'.format(set(db.collection_names()) - set(prev_collections))
            return None

        except VariantParseError, e:
            print '[ERROR] VariantParseError:', e.value
            db.drop_collection(users_variants)
            return e.value

def parse_lines(handle, file_format):
    """
    23andme:
    # rsid chromosome position genotype

    tmmb:
    chromosome position rsid reference genotype1 genotype2
    """

    ref_genome_version = None
    parse_maps = {'andme' : {'header_chr': '#',
                             'header_starts': '# rsid',
                             'fields': [('rs', 'rsid', _rs),
                                        ('chrom', 'chromosome', _chrom),
                                        ('pos', 'position', _integer),
                                        ('genotype', 'genotype', _string)],
                             'delimiter': '\t'},
                  'tmmb' : {'header_chr': '',
                            'header_starts': '',
                            'fields': [('chrom', 'chromosome', _chrom),
                                       ('pos', 'position', _integer),
                                       ('rs', 'rsid', _rs),
                                       ('ref', 'reference', _string),
                                       ('genotype1', 'genotype1', _string),
                                       ('genotype2', 'genotype2', _string),
                                       ('genotype', 'genotype', _string)],
                            'delimiter': '\t'}
                  }
    parse_map = parse_maps[file_format]

    # parse headers
    if parse_map['header_chr']:
        for line in handle:
            if not line.startswith(parse_map['header_chr']):
                raise ParseError, 'Uploaded file has no header lines. File format correct?'

            elif line.startswith(parse_map['header_starts']):
                break

            # TODO: infer ref_genome_version
            if file_format == 'andme':
                if 'build 36' in line:
                    ref_genome_version = 'b36'
                elif 'build 37' in line:
                    ref_genome_version = 'b37'

    # Parse record lines by fieldnames
    fieldnames = [x[1] for x in parse_map['fields']]
    for record in csv.DictReader(handle, fieldnames=fieldnames, delimiter=parse_map['delimiter']):
        data = {}
        for dict_name, record_name, converter in parse_map['fields']:
            try:
                data[dict_name] = converter(record.get(record_name, None))
            except VariantParseError:
                data[dict_name] = None
                data['info'] = record.get(record_name, None)

        if file_format == 'tmmb':
            data['genotype'] = data['genotype1'] + data['genotype2']

        yield data

def _rs(text):
    if text:
        try:
            return  int(text.replace('rs', ''))
        except ValueError:
            raise VariantParseError, '[WARNING] rs? {}'

    return None

def _chrom(text):
    chroms = [str(i+1) for i in range(22)] + ['X', 'Y', 'MT', 'M']
    if text:
        try:
            chrom= text.replace('chr', '')
            if not chrom in chroms:
                msg =  '[WARNING] chrom? {}'
                # print >>sys.stderr, colors.red(msg.format(text))
            return chrom
        except ValueError:
            msg =  '[WARNING] chrom? {}'
            # print >>sys.stderr, colors.red(msg.format(text))
            return chrom
    return None

def _integer(text):
    if text:
        try:
            return int(text)
        except ValueError:
            msg =  '[WARNING] integer? {}'
            # print >>sys.stderr, colors.red(msg.format(text))
            return None
    return None

def _string(text):
    return text

def _main():
    parser = argparse.ArgumentParser(description='import variants file into mongodb')
    parser.add_argument('--file-paths', metavar='filepath', nargs='+', required=True)
    parser.add_argument('--population', required=True)
    parser.add_argument('--sex', required=True)
    parser.add_argument('--file-format', required=True)
    parser.add_argument('--user-id', required=True)
    parser.add_argument('--mongo-port', default=27017)
    args = parser.parse_args()

    for file_path in args.file_paths:
        # print colors.blue('[INFO] import {}'.format(file_path))
        import_variants(file_path, args.population, args.sex, args.file_format, args.user_id, args.mongo_port)

if __name__ == '__main__':
    _main()
<commit_msg>attempt to support VCF in import_variatns.py<commit_after>#!/usr/bin/env python2.7
# -*- coding: utf-8 -*-

import argparse
import sys
import os
import subprocess
import csv
import time
import datetime

import pymongo

# import colors

class VariantParseError(Exception):
    def __init__(self, value):
        self.value = value
    def __str__(self):
        return repr(self.value)

def import_variants(file_path, population, sex, file_format, user_id, mongo_port=27017):
    """doc"""

    # remove dots from file_name
    file_name = file_path.split('/')[-1]
    file_name_cleaned = file_name.replace('.', '').replace(' ', '')
    print >>sys.stderr, '[INFO] collection name: {}'.format(file_name_cleaned)

    with pymongo.Connection(port=mongo_port) as connection:
        db = connection['pergenie']
        users_variants = db['variants'][user_id][file_name_cleaned]
        data_info = db['data_info']

        data_info.update({'user_id': user_id, 'name': file_name_cleaned}, {"$set": {'file_name_cleaned': file_name_cleaned}}, upsert=True)

        # ensure this variants file is not imported
        if users_variants.find_one():
            db.drop_collection(users_variants)
            # print >>sys.stderr, colors.yellow('[WARNING] dropped old collection of {}'.format(file_name_cleaned))

        print >>sys.stderr, '[INFO] Countiong input lines ...',
        print >>sys.stderr, '[INFO] {}'.format(file_path)
        file_lines = int(subprocess.check_output(['wc', '-l', file_path]).split()[0])
        print >>sys.stderr, 'done. # of lines: {}'.format(file_lines)

        today = str(datetime.datetime.today()).replace('-', '/')
        info = {'user_id': user_id,
                'name': file_name_cleaned,
                'raw_name': file_name,
                'date': today,
                'population': population,
                'sex': sex,
                'file_format': file_format,
                'status': float(1.0)}

        data_info.update({'user_id': user_id, 'name': file_name_cleaned}, {"$set": info}, upsert=True)

        print >>sys.stderr,'[INFO] Start importing {0} as {1}...'.format(file_name, file_name_cleaned)
        prev_collections = db.collection_names()

        try:
            with open(file_path, 'rb') as fin:
                for i,data in enumerate(parse_lines(fin, file_format)):
                    if data['rs']:
                        users_variants.insert(data)

                    if i>0 and i%10000 == 0:
                        upload_status = int( 100 * ( i*0.9 / file_lines ) + 1 )
                        data_info.update({'user_id': user_id, 'name': file_name_cleaned}, {"$set": {'status': upload_status}})
                        print '[INFO] status: {}'.format(data_info.find({'user_id': user_id, 'name': file_name_cleaned})[0]['status'])

                print >>sys.stderr,'[INFO] ensure_index ...'
                users_variants.ensure_index('rs', unique=True)  # ok?

                data_info.update({'user_id': user_id, 'name': file_name_cleaned}, {"$set": {'status': 100}})

                print >>sys.stderr, '[INFO] status: {}'.format(data_info.find({'user_id': user_id, 'name': file_name_cleaned})[0]['status'])
                print >>sys.stderr, '[INFO] done. added collection {}'.format(set(db.collection_names()) - set(prev_collections))
            return None

        except VariantParseError, e:
            print '[ERROR] VariantParseError:', e.value
            db.drop_collection(users_variants)
            return e.value

def parse_lines(handle, file_format):
    """
    23andme:
    # rsid chromosome position genotype

    tmmb:
    chromosome position rsid reference genotype1 genotype2
    """

    ref_genome_version = None
    parse_maps = {'andme' : {'header_chr': '#',
                             'header_starts': '# rsid',
                             'fields': [('rs', 'rsid', _rs),
                                        ('chrom', 'chromosome', _chrom),
                                        ('pos', 'position', _integer),
                                        ('genotype', 'genotype', _string)],
                             'delimiter': '\t'},
                  'tmmb' : {'header_chr': '',
                            'header_starts': '',
                            'fields': [('chrom', 'chromosome', _chrom),
                                       ('pos', 'position', _integer),
                                       ('rs', 'rsid', _rs),
                                       ('ref', 'reference', _string),
                                       ('genotype1', 'genotype1', _string),
                                       ('genotype2', 'genotype2', _string),
                                       ('genotype', 'genotype', _string)],
                            'delimiter': '\t',
                  'vcf' : {'header_chr': '#',
                           'header_starts': '#CHROM',
                            'fields': [('chrom', '#CHROM', _chrom),
                                       ('pos', 'POS', _integer),
                                       ('rs', 'ID', _rs),
                                       ('ref', 'REF', _string),
                                       ('alt', 'ALT', _string),
                                       ('genotype', 'genotype', _string)],
                            'delimiter': '\t',
                            }
                  }
    parse_map = parse_maps[file_format]

    # parse headers
    if parse_map['header_chr']:
        for line in handle:
            if not line.startswith(parse_map['header_chr']):
                raise ParseError, 'Uploaded file has no header lines. File format correct?'

            elif line.startswith(parse_map['header_starts']):
                fieldnames = line

                if file_format == 'vcf':
                    # get name of 1st sample (=individual)
                    vcf_first_sample_name = line.split(parse_map['vcf']['delimiter'])[9]  # ok?
                    parse_map['vcf']['fields'].append(('alt_count',
                                                       vcf_first_sample_name,
                                                       _vcf_sample2allele_count))
                break

            # TODO: infer ref_genome_version
            if file_format == 'andme':
                if 'build 36' in line:
                    ref_genome_version = 'b36'
                elif 'build 37' in line:
                    ref_genome_version = 'b37'

    # Parse record lines by fieldnames
    # fieldnames = [x[1] for x in parse_map['fields']]
    for record in csv.DictReader(handle, fieldnames=fieldnames, delimiter=parse_map['delimiter']):
        data = {}
        for dict_name, record_name, converter in parse_map['fields']:
            try:
                data[dict_name] = converter(record.get(record_name, None))
            except VariantParseError:
                data[dict_name] = None
                data['info'] = record.get(record_name, None)

        if file_format == 'tmmb':
            data['genotype'] = data['genotype1'] + data['genotype2']
        elif file_format == 'vcf':
            data['genotype'] = data['ref'] * (2 - data['alt_count']) + data['alt'] * data['alt_count']

        yield data

def _vcf_sample2alt_count(text):
    """Parse `sample` column in VCF, then retrun counts of ALT allele.
    """

def _rs(text):
    if text:
        try:
            return  int(text.replace('rs', ''))
        except ValueError:
            raise VariantParseError, '[WARNING] rs? {}'

    return None

def _chrom(text):
    chroms = [str(i+1) for i in range(22)] + ['X', 'Y', 'MT', 'M']
    if text:
        try:
            chrom= text.replace('chr', '')
            if not chrom in chroms:
                msg =  '[WARNING] chrom? {}'
                # print >>sys.stderr, colors.red(msg.format(text))
            return chrom
        except ValueError:
            msg =  '[WARNING] chrom? {}'
            # print >>sys.stderr, colors.red(msg.format(text))
            return chrom
    return None

def _integer(text):
    if text:
        try:
            return int(text)
        except ValueError:
            msg =  '[WARNING] integer? {}'
            # print >>sys.stderr, colors.red(msg.format(text))
            return None
    return None

def _string(text):
    return text

def _main():
    parser = argparse.ArgumentParser(description='import variants file into mongodb')
    parser.add_argument('--file-paths', metavar='filepath', nargs='+', required=True)
    parser.add_argument('--population', required=True)
    parser.add_argument('--sex', required=True)
    parser.add_argument('--file-format', required=True)
    parser.add_argument('--user-id', required=True)
    parser.add_argument('--mongo-port', default=27017)
    args = parser.parse_args()

    for file_path in args.file_paths:
        # print colors.blue('[INFO] import {}'.format(file_path))
        import_variants(file_path, args.population, args.sex, args.file_format, args.user_id, args.mongo_port)

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


<commit_before>__version__ = "2.1.0"

<commit_msg>Bump to v2.2.0<commit_after>__version__ = "2.2.0"

<|endoftext|>


<commit_before># Copyright (c) 2015, The MITRE Corporation. All rights reserved.
# See LICENSE.txt for complete terms.

# internal
from ramrod import utils
from ramrod.options import DEFAULT_UPDATE_OPTIONS

# relative
from . import register_updater
from .base import BaseSTIXUpdater, STIXVocab

class DiscoveryMethodVocab(STIXVocab):
    OLD_TYPES = ("DiscoveryMethodVocab-1.0",)
    NEW_TYPE = "DiscoveryMethodVocab-2.0"
    VOCAB_NAME = "STIX Default Discovery Method Vocabulary"
    VOCAB_REFERENCE = "http://stix.mitre.org/XMLSchema/default_vocabularies/1.2/stix_default_vocabularies.xsd#DiscoveryMethodTypeVocab-2.0"
    TERMS = {
       'Fraud Detection': 'External - Fraud Detection'
    }

@register_updater
class STIX_1_1_1_Updater(BaseSTIXUpdater):
    """Updates STIX v1.1.1 content to STIX v1.2.

    The following update operations are performed:

    * Instances of ``DiscoveryMethodTypeVocab-1.0`` are upgraded to
      ``DiscoveryMethodTypeVocab-2.0.``

    Note:
        All fields can be translated between STIX v1.1.1 and STIX v1.2.

    """
    VERSION = '1.1.1'

    NSMAP = {
        'TOUMarking': 'http://data-marking.mitre.org/extensions/MarkingStructure#Terms_Of_Use-1',
        'campaign': 'http://stix.mitre.org/Campaign-1',
        'coa': 'http://stix.mitre.org/CourseOfAction-1',
        'et': 'http://stix.mitre.org/ExploitTarget-1',
        'genericStructuredCOA': 'http://stix.mitre.org/extensions/StructuredCOA#Generic-1',
        'genericTM': 'http://stix.mitre.org/extensions/TestMechanism#Generic-1',
        'incident': 'http://stix.mitre.org/Incident-1',
        'indicator': 'http://stix.mitre.org/Indicator-2',
        'marking': 'http://data-marking.mitre.org/Marking-1',
        'simpleMarking': 'http://data-marking.mitre.org/extensions/MarkingStructure#Simple-1',
        'snortTM': 'http://stix.mitre.org/extensions/TestMechanism#Snort-1',
        'stix': 'http://stix.mitre.org/stix-1',
        'stix-capec': 'http://stix.mitre.org/extensions/AP#CAPEC2.7-1',
        'stix-ciqaddress': 'http://stix.mitre.org/extensions/Address#CIQAddress3.0-1',
        'stix-stix-ciq': 'http://stix.mitre.org/extensions/Identity#stix-ciq3.0-1',
        'stix-cvrf': 'http://stix.mitre.org/extensions/Vulnerability#CVRF-1',
        'stix-maec': 'http://stix.mitre.org/extensions/Malware#MAEC4.1-1',
        'stix-openioc': 'http://stix.mitre.org/extensions/TestMechanism#OpenIOC2010-1',
        'stix-oval': 'http://stix.mitre.org/extensions/TestMechanism#OVAL5.10-1',
        'stixCommon': 'http://stix.mitre.org/common-1',
        'stixVocabs': 'http://stix.mitre.org/default_vocabularies-1',
        'ta': 'http://stix.mitre.org/ThreatActor-1',
        'tlpMarking': 'http://data-marking.mitre.org/extensions/MarkingStructure#TLP-1',
        'ttp': 'http://stix.mitre.org/TTP-1',
        'yaraTM': 'http://stix.mitre.org/extensions/TestMechanism#YARA-1'
    }

    UPDATE_SCHEMALOC_MAP = {
        'http://data-marking.mitre.org/Marking-1': 'http://stix.mitre.org/XMLSchema/data_marking/1.2/data_marking.xsd',
        'http://data-marking.mitre.org/extensions/MarkingStructure#Simple-1': 'http://stix.mitre.org/XMLSchema/extensions/marking/simple/1.2/simple_marking.xsd',
        'http://data-marking.mitre.org/extensions/MarkingStructure#TLP-1': 'http://stix.mitre.org/XMLSchema/extensions/marking/tlp/1.2/tlp_marking.xsd',
        'http://data-marking.mitre.org/extensions/MarkingStructure#Terms_Of_Use-1': 'http://stix.mitre.org/XMLSchema/extensions/marking/terms_of_use/1.0.1/terms_of_use_marking.xsd',
        'http://stix.mitre.org/Campaign-1': 'http://stix.mitre.org/XMLSchema/campaign/1.2/campaign.xsd',
        'http://stix.mitre.org/CourseOfAction-1': 'http://stix.mitre.org/XMLSchema/course_of_action/1.2/course_of_action.xsd',
        'http://stix.mitre.org/ExploitTarget-1': 'http://stix.mitre.org/XMLSchema/exploit_target/1.2/exploit_target.xsd',
        'http://stix.mitre.org/Incident-1': 'http://stix.mitre.org/XMLSchema/incident/1.2/incident.xsd',
        'http://stix.mitre.org/Indicator-2': 'http://stix.mitre.org/XMLSchema/indicator/2.2/indicator.xsd',
        'http://stix.mitre.org/TTP-1': 'http://stix.mitre.org/XMLSchema/ttp/1.2/ttp.xsd',
        'http://stix.mitre.org/ThreatActor-1': 'http://stix.mitre.org/XMLSchema/threat_actor/1.2/threat_actor.xsd',
        'http://stix.mitre.org/common-1': 'http://stix.mitre.org/XMLSchema/common/1.2/stix_common.xsd',
        'http://stix.mitre.org/default_vocabularies-1': 'http://stix.mitre.org/XMLSchema/default_vocabularies/1.2.0/stix_default_vocabularies.xsd',
        'http://stix.mitre.org/extensions/AP#CAPEC2.7-1': 'http://stix.mitre.org/XMLSchema/extensions/attack_pattern/capec_2.7/1.1/capec_2.7_attack_pattern.xsd',
        'http://stix.mitre.org/extensions/Address#CIQAddress3.0-1': 'http://stix.mitre.org/XMLSchema/extensions/address/ciq_3.0/1.2/ciq_3.0_address.xsd',
        'http://stix.mitre.org/extensions/Identity#CIQIdentity3.0-1': 'http://stix.mitre.org/XMLSchema/extensions/identity/ciq_3.0/1.2/ciq_3.0_identity.xsd',
        'http://stix.mitre.org/extensions/Malware#MAEC4.1-1': 'http://stix.mitre.org/XMLSchema/extensions/malware/maec_4.1/1.1/maec_4.1_malware.xsd',
        'http://stix.mitre.org/extensions/StructuredCOA#Generic-1': 'http://stix.mitre.org/XMLSchema/extensions/structured_coa/generic/1.2/generic_structured_coa.xsd',
        'http://stix.mitre.org/extensions/TestMechanism#Generic-1': 'http://stix.mitre.org/XMLSchema/extensions/test_mechanism/generic/1.2/generic_test_mechanism.xsd',
        'http://stix.mitre.org/extensions/TestMechanism#OVAL5.10-1': 'http://stix.mitre.org/XMLSchema/extensions/test_mechanism/oval_5.10/1.2/oval_5.10_test_mechanism.xsd',
        'http://stix.mitre.org/extensions/TestMechanism#OpenIOC2010-1': 'http://stix.mitre.org/XMLSchema/extensions/test_mechanism/open_ioc_2010/1.2/open_ioc_2010_test_mechanism.xsd',
        'http://stix.mitre.org/extensions/TestMechanism#Snort-1': 'http://stix.mitre.org/XMLSchema/extensions/test_mechanism/snort/1.2/snort_test_mechanism.xsd',
        'http://stix.mitre.org/extensions/TestMechanism#YARA-1': 'http://stix.mitre.org/XMLSchema/extensions/test_mechanism/yara/1.2/yara_test_mechanism.xsd',
        'http://stix.mitre.org/extensions/Vulnerability#CVRF-1': 'http://stix.mitre.org/XMLSchema/extensions/vulnerability/cvrf_1.2/cvrf_1.1_vulnerability.xsd',
        'http://stix.mitre.org/stix-1': 'http://stix.mitre.org/XMLSchema/core/1.2/stix_core.xsd'
    }

    UPDATE_VOCABS = (
        DiscoveryMethodVocab,
    )

    def __init__(self):
        super(STIX_1_1_1_Updater, self).__init__()

    def _get_disallowed(self, root, options=None):
        """There are no untranslatable fields between STIX v1.1.1 and
        STIX v1.2.

        Note:
            This assume that `root` is schema-valid

        """
        pass

    def _get_duplicates(self, root):
        """The STIX v1.1.1 schemas enforces ID uniqueness, so
        this overrides the default ``_get_duplicates()``.

        Note:
            This assumes that `root` is schema-valid.

        """
        pass

    def _update_versions(self, root):
        """Updates the versions of versioned nodes under `root` to align with
        STIX v1.1.1 versions.

        """
        nodes = self._get_versioned_nodes(root)
        for node in nodes:
            name = utils.get_localname(node)

            if name == "Indicator":
                node.attrib['version'] = '2.2'
            else:
                node.attrib['version'] = '1.2'

    def check_update(self, root, options=None):
        """Determines if the input document can be upgraded.

        Args:
            root: The XML document. This can be a filename, a file-like object,
                an instance of ``etree._Element`` or an instance of
                ``etree._ElementTree``.
            options (optional): A ``ramrod.UpdateOptions`` instance. If
                ``None``, ``ramrod.DEFAULT_UPDATE_OPTIONS`` will be used.

        Raises:
            .UnknownVersionError: If the input document does not have a
                version.
            .InvalidVersionError: If the version of the input document
                does not match the `VERSION` class-level attribute value.

        """
        root = utils.get_etree_root(root)
        options = options or DEFAULT_UPDATE_OPTIONS

        if options.check_versions:
            self._check_version(root)

    def _update(self, root, options):
        self._update_schemalocs(root)
        self._update_versions(root)

        if options.update_vocabularies:
            self._update_vocabs(root)

        return root
<commit_msg>Modified STIX_1_1_1_Updater class docstrings.<commit_after># Copyright (c) 2015, The MITRE Corporation. All rights reserved.
# See LICENSE.txt for complete terms.

# internal
from ramrod import utils
from ramrod.options import DEFAULT_UPDATE_OPTIONS

# relative
from . import register_updater
from .base import BaseSTIXUpdater, STIXVocab

class DiscoveryMethodVocab(STIXVocab):
    OLD_TYPES = ("DiscoveryMethodVocab-1.0",)
    NEW_TYPE = "DiscoveryMethodVocab-2.0"
    VOCAB_NAME = "STIX Default Discovery Method Vocabulary"
    VOCAB_REFERENCE = "http://stix.mitre.org/XMLSchema/default_vocabularies/1.2/stix_default_vocabularies.xsd#DiscoveryMethodTypeVocab-2.0"
    TERMS = {
       'Fraud Detection': 'External - Fraud Detection'
    }

@register_updater
class STIX_1_1_1_Updater(BaseSTIXUpdater):
    """Updates STIX v1.1.1 content to STIX v1.2.

    The following update operations are performed:

    * Instances of ``DiscoveryMethodTypeVocab-1.0`` are upgraded to
      ``DiscoveryMethodTypeVocab-2.0.``
    * Component versions are updated.
    * Schemalocations are updated.

    """
    VERSION = '1.1.1'

    NSMAP = {
        'TOUMarking': 'http://data-marking.mitre.org/extensions/MarkingStructure#Terms_Of_Use-1',
        'campaign': 'http://stix.mitre.org/Campaign-1',
        'coa': 'http://stix.mitre.org/CourseOfAction-1',
        'et': 'http://stix.mitre.org/ExploitTarget-1',
        'genericStructuredCOA': 'http://stix.mitre.org/extensions/StructuredCOA#Generic-1',
        'genericTM': 'http://stix.mitre.org/extensions/TestMechanism#Generic-1',
        'incident': 'http://stix.mitre.org/Incident-1',
        'indicator': 'http://stix.mitre.org/Indicator-2',
        'marking': 'http://data-marking.mitre.org/Marking-1',
        'simpleMarking': 'http://data-marking.mitre.org/extensions/MarkingStructure#Simple-1',
        'snortTM': 'http://stix.mitre.org/extensions/TestMechanism#Snort-1',
        'stix': 'http://stix.mitre.org/stix-1',
        'stix-capec': 'http://stix.mitre.org/extensions/AP#CAPEC2.7-1',
        'stix-ciqaddress': 'http://stix.mitre.org/extensions/Address#CIQAddress3.0-1',
        'stix-stix-ciq': 'http://stix.mitre.org/extensions/Identity#stix-ciq3.0-1',
        'stix-cvrf': 'http://stix.mitre.org/extensions/Vulnerability#CVRF-1',
        'stix-maec': 'http://stix.mitre.org/extensions/Malware#MAEC4.1-1',
        'stix-openioc': 'http://stix.mitre.org/extensions/TestMechanism#OpenIOC2010-1',
        'stix-oval': 'http://stix.mitre.org/extensions/TestMechanism#OVAL5.10-1',
        'stixCommon': 'http://stix.mitre.org/common-1',
        'stixVocabs': 'http://stix.mitre.org/default_vocabularies-1',
        'ta': 'http://stix.mitre.org/ThreatActor-1',
        'tlpMarking': 'http://data-marking.mitre.org/extensions/MarkingStructure#TLP-1',
        'ttp': 'http://stix.mitre.org/TTP-1',
        'yaraTM': 'http://stix.mitre.org/extensions/TestMechanism#YARA-1'
    }

    UPDATE_SCHEMALOC_MAP = {
        'http://data-marking.mitre.org/Marking-1': 'http://stix.mitre.org/XMLSchema/data_marking/1.2/data_marking.xsd',
        'http://data-marking.mitre.org/extensions/MarkingStructure#Simple-1': 'http://stix.mitre.org/XMLSchema/extensions/marking/simple/1.2/simple_marking.xsd',
        'http://data-marking.mitre.org/extensions/MarkingStructure#TLP-1': 'http://stix.mitre.org/XMLSchema/extensions/marking/tlp/1.2/tlp_marking.xsd',
        'http://data-marking.mitre.org/extensions/MarkingStructure#Terms_Of_Use-1': 'http://stix.mitre.org/XMLSchema/extensions/marking/terms_of_use/1.0.1/terms_of_use_marking.xsd',
        'http://stix.mitre.org/Campaign-1': 'http://stix.mitre.org/XMLSchema/campaign/1.2/campaign.xsd',
        'http://stix.mitre.org/CourseOfAction-1': 'http://stix.mitre.org/XMLSchema/course_of_action/1.2/course_of_action.xsd',
        'http://stix.mitre.org/ExploitTarget-1': 'http://stix.mitre.org/XMLSchema/exploit_target/1.2/exploit_target.xsd',
        'http://stix.mitre.org/Incident-1': 'http://stix.mitre.org/XMLSchema/incident/1.2/incident.xsd',
        'http://stix.mitre.org/Indicator-2': 'http://stix.mitre.org/XMLSchema/indicator/2.2/indicator.xsd',
        'http://stix.mitre.org/TTP-1': 'http://stix.mitre.org/XMLSchema/ttp/1.2/ttp.xsd',
        'http://stix.mitre.org/ThreatActor-1': 'http://stix.mitre.org/XMLSchema/threat_actor/1.2/threat_actor.xsd',
        'http://stix.mitre.org/common-1': 'http://stix.mitre.org/XMLSchema/common/1.2/stix_common.xsd',
        'http://stix.mitre.org/default_vocabularies-1': 'http://stix.mitre.org/XMLSchema/default_vocabularies/1.2.0/stix_default_vocabularies.xsd',
        'http://stix.mitre.org/extensions/AP#CAPEC2.7-1': 'http://stix.mitre.org/XMLSchema/extensions/attack_pattern/capec_2.7/1.1/capec_2.7_attack_pattern.xsd',
        'http://stix.mitre.org/extensions/Address#CIQAddress3.0-1': 'http://stix.mitre.org/XMLSchema/extensions/address/ciq_3.0/1.2/ciq_3.0_address.xsd',
        'http://stix.mitre.org/extensions/Identity#CIQIdentity3.0-1': 'http://stix.mitre.org/XMLSchema/extensions/identity/ciq_3.0/1.2/ciq_3.0_identity.xsd',
        'http://stix.mitre.org/extensions/Malware#MAEC4.1-1': 'http://stix.mitre.org/XMLSchema/extensions/malware/maec_4.1/1.1/maec_4.1_malware.xsd',
        'http://stix.mitre.org/extensions/StructuredCOA#Generic-1': 'http://stix.mitre.org/XMLSchema/extensions/structured_coa/generic/1.2/generic_structured_coa.xsd',
        'http://stix.mitre.org/extensions/TestMechanism#Generic-1': 'http://stix.mitre.org/XMLSchema/extensions/test_mechanism/generic/1.2/generic_test_mechanism.xsd',
        'http://stix.mitre.org/extensions/TestMechanism#OVAL5.10-1': 'http://stix.mitre.org/XMLSchema/extensions/test_mechanism/oval_5.10/1.2/oval_5.10_test_mechanism.xsd',
        'http://stix.mitre.org/extensions/TestMechanism#OpenIOC2010-1': 'http://stix.mitre.org/XMLSchema/extensions/test_mechanism/open_ioc_2010/1.2/open_ioc_2010_test_mechanism.xsd',
        'http://stix.mitre.org/extensions/TestMechanism#Snort-1': 'http://stix.mitre.org/XMLSchema/extensions/test_mechanism/snort/1.2/snort_test_mechanism.xsd',
        'http://stix.mitre.org/extensions/TestMechanism#YARA-1': 'http://stix.mitre.org/XMLSchema/extensions/test_mechanism/yara/1.2/yara_test_mechanism.xsd',
        'http://stix.mitre.org/extensions/Vulnerability#CVRF-1': 'http://stix.mitre.org/XMLSchema/extensions/vulnerability/cvrf_1.2/cvrf_1.1_vulnerability.xsd',
        'http://stix.mitre.org/stix-1': 'http://stix.mitre.org/XMLSchema/core/1.2/stix_core.xsd'
    }

    UPDATE_VOCABS = (
        DiscoveryMethodVocab,
    )

    def __init__(self):
        super(STIX_1_1_1_Updater, self).__init__()

    def _get_disallowed(self, root, options=None):
        """There are no untranslatable fields between STIX v1.1.1 and
        STIX v1.2.

        Note:
            This assume that `root` is schema-valid

        """
        pass

    def _get_duplicates(self, root):
        """The STIX v1.1.1 schemas enforces ID uniqueness, so
        this overrides the default ``_get_duplicates()``.

        Note:
            This assumes that `root` is schema-valid.

        """
        pass

    def _update_versions(self, root):
        """Updates the versions of versioned nodes under `root` to align with
        STIX v1.1.1 versions.

        """
        nodes = self._get_versioned_nodes(root)
        for node in nodes:
            name = utils.get_localname(node)

            if name == "Indicator":
                node.attrib['version'] = '2.2'
            else:
                node.attrib['version'] = '1.2'

    def check_update(self, root, options=None):
        """Determines if the input document can be upgraded.

        Args:
            root: The XML document. This can be a filename, a file-like object,
                an instance of ``etree._Element`` or an instance of
                ``etree._ElementTree``.
            options (optional): A ``ramrod.UpdateOptions`` instance. If
                ``None``, ``ramrod.DEFAULT_UPDATE_OPTIONS`` will be used.

        Raises:
            .UnknownVersionError: If the input document does not have a
                version.
            .InvalidVersionError: If the version of the input document
                does not match the `VERSION` class-level attribute value.

        """
        root = utils.get_etree_root(root)
        options = options or DEFAULT_UPDATE_OPTIONS

        if options.check_versions:
            self._check_version(root)

    def _update(self, root, options):
        self._update_schemalocs(root)
        self._update_versions(root)

        if options.update_vocabularies:
            self._update_vocabs(root)

        return root
<|endoftext|>


<commit_before># Copyright (c) 2013 Per Unneberg
#
# Licensed under the Apache License, Version 2.0 (the "License"); you may not
# use this file except in compliance with the License. You may obtain a copy of
# the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
# WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
# License for the specific language governing permissions and limitations under
# the License.
"""
Wrapper library for `bwa <http://bio-bwa.sourceforge.net/>`_.

Classes
--------
"""
import re
import luigi
from ratatosk.job import InputJobTask, JobTask, JobWrapperTask, DefaultShellJobRunner, PipedTask
from ratatosk.lib.tools.samtools import SamToBam
from ratatosk.utils import rreplace, fullclassname

class BwaJobRunner(DefaultShellJobRunner):
    pass

class InputFastqFile(InputJobTask):
    _config_section = "bwa"
    _config_subsection = "InputFastqFile"
    parent_task = luigi.Parameter(default="ratatosk.lib.files.external.FastqFile")
    suffix = luigi.Parameter(default=".fastq.gz")

class InputFastaFile(InputJobTask):
    _config_section = "bwa"
    _config_subsection = "InputFastaFile"
    parent_task = luigi.Parameter(default="ratatosk.lib.files.external.FastaFile")
    suffix = luigi.Parameter(default=".fa")

class BwaJobTask(JobTask):
    """Main bwa class with parameters necessary for all bwa classes"""
    _config_section = "bwa"
    executable = luigi.Parameter(default="bwa")
    bwaref = luigi.Parameter(default=None)
    num_threads = luigi.Parameter(default=1)

    def job_runner(self):
        return BwaJobRunner()

    def output(self):
        return luigi.LocalTarget(self.target)

class Aln(BwaJobTask):
    _config_subsection = "Aln"
    sub_executable = "aln"
    parent_task = luigi.Parameter(default=("ratatosk.lib.align.bwa.InputFastqFile",))
    suffix = luigi.Parameter(default=".sai")
    read1_suffix = luigi.Parameter(default="_R1_001")
    read2_suffix = luigi.Parameter(default="_R2_001")
    is_read1 = True
    can_multi_thread = True

    def opts(self):
        # Threads is an option so handle it here
        retval = list(self.options)
        return retval + ['-t {}'.format(str(self.threads()))]

    def requires(self):
        cls = self.parent()[0]
        source = self.source()[0]
        # Ugly hack for 1 -> 2 dependency: works but should be dealt with otherwise
        if str(fullclassname(cls)) in ["ratatosk.lib.utils.misc.ResyncMatesJobTask"]:
            if re.search(self.read1_suffix, source):
                self.is_read1 = True
                fq1 = source
                fq2 = rreplace(source, self.read1_suffix, self.read2_suffix, 1)
            else:
                self.is_read1 = False
                fq1 = rreplace(source, self.read2_suffix, self.read1_suffix, 1)
                fq2 = source
            retval = [cls(target=[fq1, fq2])]
        else:
            retval = [cls(target=source)]
        if len(self.parent()) > 1:
            retval += [cls(target=source) for cls in self.parent()[1:]]
        return retval

    def args(self):
        # bwa aln "-f" option seems to be broken!?!
        if isinstance(self.input(), list):
            if self.is_read1:
                return [self.bwaref, self.input()[0], ">", self.output()]
            else:
                return [self.bwaref, self.input()[1], ">", self.output()]
        else:
            return [self.bwaref, self.input(), ">", self.output()]

class BwaAlnWrapperTask(JobWrapperTask):
    fastqfiles = luigi.Parameter(default=[], is_list=True)
    def requires(self):
        return [Aln(target=x) for x in self.fastqfiles]

class Sampe(BwaJobTask):
    _config_subsection = "Sampe"
    sub_executable = "sampe"
    add_label = luigi.Parameter(default=("_R1_001", "_R2_001"), is_list=True)
    suffix = luigi.Parameter(default=".sam")
    read_group = luigi.Parameter(default=None)
    platform = luigi.Parameter(default="Illumina")
    parent_task = luigi.Parameter(default=("ratatosk.lib.align.bwa.Aln", "ratatosk.lib.align.bwa.Aln"), is_list=True)
    can_multi_thread = False
    max_memory_gb = 6 # bwa documentation says ~5.4 for human genome

    def _get_read_group(self):
        if not self.read_group:
            from ratatosk import backend
            cls = self.parent()[0]
            sai1 = self.input()[0]
            rgid = rreplace(rreplace(sai1.path, cls().suffix, "", 1), self.add_label[0], "", 1)
            smid = rgid
            # Get sample information if present in global vars. Note
            # that this requires the
            # backend.__global_vars__["targets"] be set
            # This is not perfect but works for now
            for tgt in backend.__global_vars__.get("targets", []):
                if smid.startswith(tgt[2]):
                    smid = tgt[0]
                    break
            # The platform should be configured elsewhere
            rg = "\"{}\"".format("\t".join(["@RG", "ID:{}".format(rgid), "SM:{}".format(smid), "PL:{}".format(self.platform)]))
            if self.pipe:
                return rg.replace("\t", "\\t")
            else:
                return rg
        else:
            return self.read_group

    def args(self):
        cls = self.parent()[0]
        parent_cls = cls().parent()[0]
        (fastq1, fastq2) = [luigi.LocalTarget(rreplace(sai.path, cls().suffix, parent_cls().sfx(), 1)) for sai in self.input()]
        return ["-r", self._get_read_group(), self.bwaref, self.input()[0].path, self.input()[1].path, fastq1, fastq2, ">", self.output()]

class Bampe(PipedTask):
    _config_section = "bwa"
    _config_subsection = "Bampe"
    add_label = luigi.Parameter(default=("_R1_001", "_R2_001"), is_list=True)
    parent_task = luigi.Parameter(default=("ratatosk.lib.align.bwa.Aln", "ratatosk.lib.align.bwa.Aln"), is_list=True)
    suffix = luigi.Parameter(default=".bam")
    read_group = luigi.Parameter(default=None)
    platform = luigi.Parameter(default="Illumina")
    can_multi_thread = False
    max_memory_gb = 6 # bwa documentation says ~5.4 for human genome

    def args(self):
        return [Sampe(target=self.target.replace(".bam", ".sam"), pipe=True), SamToBam(target=self.target, pipe=True)]

class Index(BwaJobTask):
    _config_subsection = "index"
    sub_executable = "index"
    suffix = luigi.Parameter(default=".fa.bwt")
    parent_task = luigi.Parameter(default=("ratatosk.lib.align.bwa.InputFastaFile", ), is_list=True)

    def args(self):
        return [self.input()[0]]
<commit_msg>Update bwa aln bug fix<commit_after># Copyright (c) 2013 Per Unneberg
#
# Licensed under the Apache License, Version 2.0 (the "License"); you may not
# use this file except in compliance with the License. You may obtain a copy of
# the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
# WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
# License for the specific language governing permissions and limitations under
# the License.
"""
Wrapper library for `bwa <http://bio-bwa.sourceforge.net/>`_.

Classes
--------
"""
import re
import luigi
from itertools import izip
from ratatosk.job import InputJobTask, JobTask, JobWrapperTask, DefaultShellJobRunner, PipedTask
from ratatosk.lib.tools.samtools import SamToBam
from ratatosk.utils import rreplace, fullclassname

class BwaJobRunner(DefaultShellJobRunner):
    pass

class InputFastqFile(InputJobTask):
    _config_section = "bwa"
    _config_subsection = "InputFastqFile"
    parent_task = luigi.Parameter(default="ratatosk.lib.files.external.FastqFile")
    suffix = luigi.Parameter(default=".fastq.gz")

class InputFastaFile(InputJobTask):
    _config_section = "bwa"
    _config_subsection = "InputFastaFile"
    parent_task = luigi.Parameter(default="ratatosk.lib.files.external.FastaFile")
    suffix = luigi.Parameter(default=".fa")

class BwaJobTask(JobTask):
    """Main bwa class with parameters necessary for all bwa classes"""
    _config_section = "bwa"
    executable = luigi.Parameter(default="bwa")
    bwaref = luigi.Parameter(default=None)
    num_threads = luigi.Parameter(default=1)

    def job_runner(self):
        return BwaJobRunner()

    def output(self):
        return luigi.LocalTarget(self.target)

class Aln(BwaJobTask):
    _config_subsection = "Aln"
    sub_executable = "aln"
    parent_task = luigi.Parameter(default=("ratatosk.lib.align.bwa.InputFastqFile",))
    suffix = luigi.Parameter(default=".sai")
    read1_suffix = luigi.Parameter(default="_R1_001")
    read2_suffix = luigi.Parameter(default="_R2_001")
    is_read1 = True
    can_multi_thread = True

    def opts(self):
        # Threads is an option so handle it here
        retval = list(self.options)
        return retval + ['-t {}'.format(str(self.threads()))]

    def requires(self):
        cls = self.parent()[0]
        source = self.source()[0]
        # Ugly hack for 1 -> 2 dependency: works but should be dealt with otherwise
        if str(fullclassname(cls)) in ["ratatosk.lib.utils.misc.ResyncMatesJobTask"]:
            if re.search(self.read1_suffix, source):
                self.is_read1 = True
                fq1 = source
                fq2 = rreplace(source, self.read1_suffix, self.read2_suffix, 1)
            else:
                self.is_read1 = False
                fq1 = rreplace(source, self.read2_suffix, self.read1_suffix, 1)
                fq2 = source
            retval = [cls(target=[fq1, fq2])]
        else:
            retval = [cls(target=source)]
        if len(self.parent()) > 1:
            retval += [cls(target=source) for cls, source in izip(self.parent()[1:], self.source()[1:])]
        return retval

    def args(self):
        # bwa aln "-f" option seems to be broken!?!
        if isinstance(self.input(), list):
            if self.is_read1:
                return [self.bwaref, self.input()[0], ">", self.output()]
            else:
                return [self.bwaref, self.input()[1], ">", self.output()]
        else:
            return [self.bwaref, self.input(), ">", self.output()]

class BwaAlnWrapperTask(JobWrapperTask):
    fastqfiles = luigi.Parameter(default=[], is_list=True)
    def requires(self):
        return [Aln(target=x) for x in self.fastqfiles]

class Sampe(BwaJobTask):
    _config_subsection = "Sampe"
    sub_executable = "sampe"
    add_label = luigi.Parameter(default=("_R1_001", "_R2_001"), is_list=True)
    suffix = luigi.Parameter(default=".sam")
    read_group = luigi.Parameter(default=None)
    platform = luigi.Parameter(default="Illumina")
    parent_task = luigi.Parameter(default=("ratatosk.lib.align.bwa.Aln", "ratatosk.lib.align.bwa.Aln"), is_list=True)
    can_multi_thread = False
    max_memory_gb = 6 # bwa documentation says ~5.4 for human genome

    def _get_read_group(self):
        if not self.read_group:
            from ratatosk import backend
            cls = self.parent()[0]
            sai1 = self.input()[0]
            rgid = rreplace(rreplace(sai1.path, cls().suffix, "", 1), self.add_label[0], "", 1)
            smid = rgid
            # Get sample information if present in global vars. Note
            # that this requires the
            # backend.__global_vars__["targets"] be set
            # This is not perfect but works for now
            for tgt in backend.__global_vars__.get("targets", []):
                if smid.startswith(tgt[2]):
                    smid = tgt[0]
                    break
            # The platform should be configured elsewhere
            rg = "\"{}\"".format("\t".join(["@RG", "ID:{}".format(rgid), "SM:{}".format(smid), "PL:{}".format(self.platform)]))
            if self.pipe:
                return rg.replace("\t", "\\t")
            else:
                return rg
        else:
            return self.read_group

    def args(self):
        cls = self.parent()[0]
        parent_cls = cls().parent()[0]
        (fastq1, fastq2) = [luigi.LocalTarget(rreplace(sai.path, cls().suffix, parent_cls().sfx(), 1)) for sai in self.input()]
        return ["-r", self._get_read_group(), self.bwaref, self.input()[0].path, self.input()[1].path, fastq1, fastq2, ">", self.output()]

class Bampe(PipedTask):
    _config_section = "bwa"
    _config_subsection = "Bampe"
    add_label = luigi.Parameter(default=("_R1_001", "_R2_001"), is_list=True)
    parent_task = luigi.Parameter(default=("ratatosk.lib.align.bwa.Aln", "ratatosk.lib.align.bwa.Aln"), is_list=True)
    suffix = luigi.Parameter(default=".bam")
    read_group = luigi.Parameter(default=None)
    platform = luigi.Parameter(default="Illumina")
    can_multi_thread = False
    max_memory_gb = 6 # bwa documentation says ~5.4 for human genome

    def args(self):
        return [Sampe(target=self.target.replace(".bam", ".sam"), pipe=True), SamToBam(target=self.target, pipe=True)]

class Index(BwaJobTask):
    _config_subsection = "index"
    sub_executable = "index"
    suffix = luigi.Parameter(default=".fa.bwt")
    parent_task = luigi.Parameter(default=("ratatosk.lib.align.bwa.InputFastaFile", ), is_list=True)

    def args(self):
        return [self.input()[0]]
<|endoftext|>


<commit_before>
#
# spyne - Copyright (C) Spyne contributors.
#
# This library is free software; you can redistribute it and/or
# modify it under the terms of the GNU Lesser General Public
# License as published by the Free Software Foundation; either
# version 2.1 of the License, or (at your option) any later version.
#
# This library is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
# Lesser General Public License for more details.
#
# You should have received a copy of the GNU Lesser General Public
# License along with this library; if not, write to the Free Software
# Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301
#

class LogicError(Exception):
    pass

__version__ = '2.13.12-alpha'

from pytz import utc as LOCAL_TZ
from decimal import Decimal as D

DEFAULT_LANGUAGE = 'en'

from spyne._base import Address

from spyne.context import AuxMethodContext
from spyne.context import TransportContext
from spyne.context import ProtocolContext
from spyne.context import EventContext
from spyne.context import MethodContext

from spyne.evmgr import EventManager

from spyne.descriptor import MethodDescriptor
from spyne.descriptor import BODY_STYLE_WRAPPED
from spyne.descriptor import BODY_STYLE_BARE
from spyne.descriptor import BODY_STYLE_OUT_BARE
from spyne.descriptor import BODY_STYLE_EMPTY
from spyne.descriptor import BODY_STYLE_EMPTY_OUT_BARE

# decorator imports descriptor, so this needs to come after
from spyne.decorator import rpc
from spyne.decorator import srpc
from spyne.decorator import mrpc

from spyne.service import ServiceBase as Service
from spyne.service import ServiceBase  # DEPRECATED
from spyne.application import Application

from spyne.model import *
from spyne.model import Mandatory as M

from spyne.error import InvalidCredentialsError
from spyne.error import RequestTooLongError
from spyne.error import RequestNotAllowed
from spyne.error import ArgumentError
from spyne.error import InvalidInputError
from spyne.error import MissingFieldError
from spyne.error import ValidationError
from spyne.error import InternalError
from spyne.error import ResourceNotFoundError
from spyne.error import RespawnError
from spyne.error import ResourceAlreadyExistsError
from spyne.error import Redirect

from spyne.client import ClientBase, RemoteProcedureBase, RemoteService
from spyne.server import ServerBase, NullServer

def _vercheck():
    import sys
    if not hasattr(sys, "version_info") or sys.version_info < (2, 6):
        raise RuntimeError("Spyne requires Python 2.6 or later. Trust us.")
_vercheck()
<commit_msg>verbump<commit_after>
#
# spyne - Copyright (C) Spyne contributors.
#
# This library is free software; you can redistribute it and/or
# modify it under the terms of the GNU Lesser General Public
# License as published by the Free Software Foundation; either
# version 2.1 of the License, or (at your option) any later version.
#
# This library is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
# Lesser General Public License for more details.
#
# You should have received a copy of the GNU Lesser General Public
# License along with this library; if not, write to the Free Software
# Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301
#

class LogicError(Exception):
    pass

__version__ = '2.13.12-beta'

from pytz import utc as LOCAL_TZ
from decimal import Decimal as D

DEFAULT_LANGUAGE = 'en'

from spyne._base import Address

from spyne.context import AuxMethodContext
from spyne.context import TransportContext
from spyne.context import ProtocolContext
from spyne.context import EventContext
from spyne.context import MethodContext

from spyne.evmgr import EventManager

from spyne.descriptor import MethodDescriptor
from spyne.descriptor import BODY_STYLE_WRAPPED
from spyne.descriptor import BODY_STYLE_BARE
from spyne.descriptor import BODY_STYLE_OUT_BARE
from spyne.descriptor import BODY_STYLE_EMPTY
from spyne.descriptor import BODY_STYLE_EMPTY_OUT_BARE

# decorator imports descriptor, so this needs to come after
from spyne.decorator import rpc
from spyne.decorator import srpc
from spyne.decorator import mrpc

from spyne.service import ServiceBase as Service
from spyne.service import ServiceBase  # DEPRECATED
from spyne.application import Application

from spyne.model import *
from spyne.model import Mandatory as M

from spyne.error import InvalidCredentialsError
from spyne.error import RequestTooLongError
from spyne.error import RequestNotAllowed
from spyne.error import ArgumentError
from spyne.error import InvalidInputError
from spyne.error import MissingFieldError
from spyne.error import ValidationError
from spyne.error import InternalError
from spyne.error import ResourceNotFoundError
from spyne.error import RespawnError
from spyne.error import ResourceAlreadyExistsError
from spyne.error import Redirect

from spyne.client import ClientBase, RemoteProcedureBase, RemoteService
from spyne.server import ServerBase, NullServer

def _vercheck():
    import sys
    if not hasattr(sys, "version_info") or sys.version_info < (2, 6):
        raise RuntimeError("Spyne requires Python 2.6 or later. Trust us.")
_vercheck()
<|endoftext|>


<commit_before>
from __future__ import print_function, unicode_literals
import argparse
import codecs
import sys

if sys.version_info.major == 2:
    sys.stdin = codecs.getreader('utf-8')(sys.stdin)
    sys.stdout = codecs.getwriter('utf-8')(sys.stdout)

from depccg import PyAStarParser, PyJaAStarParser, to_mathml

Parsers = {"en": PyAStarParser, "ja": PyJaAStarParser}

class Token:
    def __init__(self, word, lemma, pos, chunk, entity):
        self.word   = word
        self.lemma  = lemma
        self.pos    = pos
        self.chunk  = chunk
        self.entity = entity

    @staticmethod
    def from_piped(string):
        # WORD|POS|NER or WORD|LEMMA|POS|NER
        # or WORD|LEMMA|POS|NER|CHUCK
        items = string.split("|")
        if len(items) == 5:
            w, l, p, n, c = items
            return Token(w, l, p, c, n)
        elif len(items) == 4:
            w, l, p, n = items
            return Token(w, l, p, "XX", n)
        else:
            w, p, n = items
            return Token(w, "XX", p, "XX", n)

def to_xml(trees, tagged_doc, file=sys.stdout):
    print("<?xml version=\"1.0\" encoding=\"UTF-8\"?>")
    print("<?xml-stylesheet type=\"text/xsl\" href=\"candc.xml\"?>")
    print("<candc>")
    for i, (tree, tagged) in enumerate(zip(trees, tagged_doc), 1):
        for j, (t, _) in enumerate(tree, 1):
            print("<ccg sentence=\"{}\" id=\"{}\">".format(i, j))
            print(t.xml.format(*tagged))
            print("</ccg>")
    print("</candc>")

def to_prolog(trees, tagged_doc, file=sys.stdout):
    print(":- op(601, xfx, (/)).")
    print(":- op(601, xfx, (\)).")
    print(":- multifile ccg/2, id/2.")
    print(":- discontiguous ccg/2, id/2.\n")
    for i, (tree, tagged) in enumerate(zip(trees, tagged_doc), 1):
        for tok in tagged:
            if "'" in tok.lemma:
                tok.lemma = tok.lemma.replace("'", "\\'")
        for t, _ in tree:
            print(t.prolog.format(i, *tagged))

if __name__ == "__main__":
    parser = argparse.ArgumentParser("A* CCG parser")
    parser.add_argument("model", help="model directory")
    parser.add_argument("lang", help="language", choices=["en", "ja"])
    parser.add_argument("--input", default=None,
            help="a file with tokenized sentences in each line")
    parser.add_argument("--batchsize", type=int, default=32,
            help="batchsize in supertagger")
    parser.add_argument("--nbest", type=int, default=1,
            help="output N best parses")
    parser.add_argument("--input-format", default="raw",
            choices=["raw", "POSandNERtagged"],
            help="input format")
    parser.add_argument("--format", default="auto",
            choices=["auto", "deriv", "xml", "ja", "conll", "html", "prolog"],
            help="output format")
    parser.add_argument("--root-cats", default=None,
            help="allow only these categories to be at the root of a tree. If None, use default setting.")

    parser.add_argument("--verbose", action="store_true")
    args = parser.parse_args()

    fin = sys.stdin if args.input is None else codecs.open(args.input, encoding="utf-8")
    if args.input_format == "POSandNERtagged":
        tagged_doc = [[Token.from_piped(token) for token in sent.strip().split(" ")]for sent in fin]
        doc = [" ".join([token.word for token in sent]) for sent in tagged_doc]
    else:
        assert args.format not in ["xml", "prolog"], \
                "XML and Prolog output format is supported only with --input-format POSandNERtagged."
        doc = [l.strip() for l in fin]

    if args.root_cats is not None:
        args.root_cats = args.root_cats.split(",")

    parser = Parsers[args.lang](args.model,
                               batchsize=args.batchsize,
                               nbest=args.nbest,
                               possible_root_cats=args.root_cats,
                               loglevel=1 if args.verbose else 3)
    res = parser.parse_doc(doc)

    if args.format == "xml":
        to_xml(res, tagged_doc)
    elif args.format == "prolog":
        to_prolog(res, tagged_doc)
    elif args.format == "html":
        to_mathml(res, file=sys.stdout)
    elif args.format == "conll":
        for i, parsed in enumerate(res):
            for tree, prob in parsed:
                print("# ID={}\n# log probability={:.4e}\n{}".format(
                        i, prob, tree.conll))
    else: # "auto", "deriv", "ja"
        for i, parsed in enumerate(res, 1):
            for tree, _ in parsed:
                print("ID={}".format(i))
                print(getattr(tree, args.format))

<commit_msg>update<commit_after>
from __future__ import print_function, unicode_literals
import argparse
import codecs
import sys

try:
    import html
    def escape(text):
        return html.escape(text)
except ImportError:
    import cgi
    def escape(text):
        return cgi.escape(text)

if sys.version_info.major == 2:
    sys.stdin = codecs.getreader('utf-8')(sys.stdin)
    sys.stdout = codecs.getwriter('utf-8')(sys.stdout)

from depccg import PyAStarParser, PyJaAStarParser, to_mathml

Parsers = {"en": PyAStarParser, "ja": PyJaAStarParser}

class Token:
    def __init__(self, word, lemma, pos, chunk, entity):
        self.word   = word
        self.lemma  = lemma
        self.pos    = pos
        self.chunk  = chunk
        self.entity = entity

    @staticmethod
    def from_piped(string):
        # WORD|POS|NER or WORD|LEMMA|POS|NER
        # or WORD|LEMMA|POS|NER|CHUCK
        items = string.split("|")
        if len(items) == 5:
            w, l, p, n, c = items
            return Token(w, l, p, c, n)
        elif len(items) == 4:
            w, l, p, n = items
            return Token(w, l, p, "XX", n)
        else:
            w, p, n = items
            return Token(w, "XX", p, "XX", n)

    def escaped(self):
        return Token(
                escape(self.word),
                escape(self.lemma),
                escape(self.pos),
                escape(self.chunk),
                escape(self.entity))

def to_xml(trees, tagged_doc, file=sys.stdout):
    print("<?xml version=\"1.0\" encoding=\"UTF-8\"?>")
    print("<?xml-stylesheet type=\"text/xsl\" href=\"candc.xml\"?>")
    print("<candc>")
    for i, (tree, tagged) in enumerate(zip(trees, tagged_doc), 1):
        for j, (t, _) in enumerate(tree, 1):
            print("<ccg sentence=\"{}\" id=\"{}\">".format(i, j))
            print(t.xml.format(*[token.escaped() for token in tagged]))
            print("</ccg>")
    print("</candc>")

def to_prolog(trees, tagged_doc, file=sys.stdout):
    print(":- op(601, xfx, (/)).")
    print(":- op(601, xfx, (\)).")
    print(":- multifile ccg/2, id/2.")
    print(":- discontiguous ccg/2, id/2.\n")
    for i, (tree, tagged) in enumerate(zip(trees, tagged_doc), 1):
        for tok in tagged:
            if "'" in tok.lemma:
                tok.lemma = tok.lemma.replace("'", "\\'")
        for t, _ in tree:
            print(t.prolog.format(i, *tagged))

if __name__ == "__main__":
    parser = argparse.ArgumentParser("A* CCG parser")
    parser.add_argument("model", help="model directory")
    parser.add_argument("lang", help="language", choices=["en", "ja"])
    parser.add_argument("--input", default=None,
            help="a file with tokenized sentences in each line")
    parser.add_argument("--batchsize", type=int, default=32,
            help="batchsize in supertagger")
    parser.add_argument("--nbest", type=int, default=1,
            help="output N best parses")
    parser.add_argument("--input-format", default="raw",
            choices=["raw", "POSandNERtagged"],
            help="input format")
    parser.add_argument("--format", default="auto",
            choices=["auto", "deriv", "xml", "ja", "conll", "html", "prolog"],
            help="output format")
    parser.add_argument("--root-cats", default=None,
            help="allow only these categories to be at the root of a tree. If None, use default setting.")

    parser.add_argument("--verbose", action="store_true")
    args = parser.parse_args()

    fin = sys.stdin if args.input is None else codecs.open(args.input, encoding="utf-8")
    if args.input_format == "POSandNERtagged":
        tagged_doc = [[Token.from_piped(token) for token in sent.strip().split(" ")]for sent in fin]
        doc = [" ".join([token.word for token in sent]) for sent in tagged_doc]
    else:
        assert args.format not in ["xml", "prolog"], \
                "XML and Prolog output format is supported only with --input-format POSandNERtagged."
        doc = [l.strip() for l in fin]

    if args.root_cats is not None:
        args.root_cats = args.root_cats.split(",")

    parser = Parsers[args.lang](args.model,
                               batchsize=args.batchsize,
                               nbest=args.nbest,
                               possible_root_cats=args.root_cats,
                               loglevel=1 if args.verbose else 3)
    res = parser.parse_doc(doc)

    if args.format == "xml":
        to_xml(res, tagged_doc)
    elif args.format == "prolog":
        to_prolog(res, tagged_doc)
    elif args.format == "html":
        to_mathml(res, file=sys.stdout)
    elif args.format == "conll":
        for i, parsed in enumerate(res):
            for tree, prob in parsed:
                print("# ID={}\n# log probability={:.4e}\n{}".format(
                        i, prob, tree.conll))
    else: # "auto", "deriv", "ja"
        for i, parsed in enumerate(res, 1):
            for tree, _ in parsed:
                print("ID={}".format(i))
                print(getattr(tree, args.format))

<|endoftext|>