<commit_before>from __future__ import unicode_literals

from bson.dbref import DBRef
from datetime import datetime
from mongoengine import *

from core.database import YetiDocument


class InvestigationLink(EmbeddedDocument):
    id = StringField(required=True)
    fromnode = StringField(required=True)
    tonode = StringField(required=True)
    label = StringField()

    @staticmethod
    def build(data):
        link = InvestigationLink(id=data['id'], fromnode=data['from'], tonode=data['to'])
        if 'label' in data:
            link.label = data['label']

        return link


class InvestigationEvent(EmbeddedDocument):
    kind = StringField(required=True)
    links = ListField(EmbeddedDocumentField(InvestigationLink))
    nodes = ListField(ReferenceField('Node'))
    datetime = DateTimeField(default=datetime.utcnow)


class Investigation(YetiDocument):
    name = StringField()
    description = StringField()
    links = ListField(EmbeddedDocumentField(InvestigationLink))
    nodes = ListField(ReferenceField('Node', dbref=True))
    events = ListField(EmbeddedDocumentField(InvestigationEvent))
    created = DateTimeField(default=datetime.utcnow)
    updated = DateTimeField(default=datetime.utcnow)

    def info(self):
        result = self.to_mongo()
        result['nodes'] = [node.to_mongo() for node in self.nodes]

        return result

    def _node_changes(self, kind, method, links, nodes):
        event = InvestigationEvent(kind=kind)

        for link in links:
            link = InvestigationLink.build(link)
            if method('links', link.to_mongo()):
                event.links.append(link)

        for node in nodes:
            if not isinstance(node, DBRef):
                node = node.to_dbref()

            if method('nodes', node):
                event.nodes.append(node)<commit_msg>Bugfix<commit_after>from __future__ import unicode_literals

from bson.dbref import DBRef
from datetime import datetime
from mongoengine import *

from core.database import YetiDocument


class InvestigationLink(EmbeddedDocument):
    id = StringField(required=True)
    fromnode = StringField(required=True)
    tonode = StringField(required=True)
    label = StringField()

    @staticmethod
    def build(data):
        link = InvestigationLink(id=data['id'], fromnode=data['from'], tonode=data['to'])
        if 'label' in data:
            link.label = data['label']

        return link


class InvestigationEvent(EmbeddedDocument):
    kind = StringField(required=True)
    links = ListField(EmbeddedDocumentField(InvestigationLink))
    nodes = ListField(ReferenceField('Node'))
    datetime = DateTimeField(default=datetime.utcnow)


class Investigation(YetiDocument):
    name = StringField()
    description = StringField()
    links = ListField(EmbeddedDocumentField(InvestigationLink))
    nodes = ListField(ReferenceField('Node', dbref=True))
    events = ListField(EmbeddedDocumentField(InvestigationEvent))
    created = DateTimeField(default=datetime.utcnow)
    updated = DateTimeField(default=datetime.utcnow)

    SEARCH_ALIASES = {}

    def info(self):
        result = self.to_mongo()
        result['nodes'] = [node.to_mongo() for node in self.nodes]

        return result

    def _node_changes(self, kind, method, links, nodes):
        event = InvestigationEvent(kind=kind)

        for link in links:
            link = InvestigationLink.build(link)
            if method('links', link.to_mongo()):
                event.links.append(link)

        for node in nodes:
            if not isinstance(node, DBRef):
                node = node.to_dbref()

            if method('nodes', node):
                event.nodes.append(node)
{'fraction_non_alphanumeric': 0.10081081081081081, 'fraction_numerical': 0.0, 'mean_word_length': 3.308498253783469, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '4187523', 'n_tokens_mistral': 1061, 'n_tokens_neox': 999, 'n_words': 230}