<commit_before>import re
import json
import time
import logging
import itertools
from ndex2.nice_cx_network import NiceCXNetwork
from collections import OrderedDict
from indra.statements import *
from indra.databases import context_client, ndex_client, get_identifiers_url


logger = logging.getLogger(__name__)


class NiceCxAssembler(object):
    def __init__(self, statements=None, name=None):
        self.statements = statements if statements else []
        self.network = NiceCXNetwork()
        self.network.set_network_attribute('name', (name if name else
                                                    'indra_assembled'))
        self.node_keys = {}
        self.citation_keys = {}

    def make_model(self):
        for stmt in self.statements:
            agents = stmt.agent_list()
            not_none_agents = [a for a in agents if a is not None]
            if len(not_none_agents) < 2:
                continue
            for a1, a2 in itertools.combinations(not_none_agents, 2):
                a1_id = self.add_node(a1)
                a2_id = self.add_node(a2)
                edge_id = self.add_edge(a1_id, a2_id, stmt)
        prefixes = {'pubmed': 'http://identifiers.org/pubmed/'}
        self.network.set_network_attribute('@context', json.dumps(prefixes))

    def add_node(self, agent):
        agent_key = self.get_agent_key(agent)
        if agent_key in self.node_keys:
            node_id = self.node_keys[agent_key]
        else:
            node_id = self.network.create_node(agent.name, str(agent.db_refs))
            self.node_keys[agent_key] = node_id
            self.network.add_node_attribute(property_of=node_id,
                                            name='db_refs',
                                            values=json.dumps(agent.db_refs),
                                            type='string')
        return node_id

    def add_edge(self, a1_id, a2_id, stmt):
        stmt_type = stmt.__class__.__name__
        edge_id = self.network.create_edge(a1_id, a2_id, stmt_type)
        pmids = [ev.pmid for ev in stmt.evidence if ev.pmid is not None]
        self.network.set_edge_attribute(edge_id, 'citation',
                                        ['pubmed:%s' % p for p in pmids],
                                        type='list_of_string')
        return edge_id

    def print_model(self):
        return self.network.to_cx()

    def get_agent_key(self, agent):
        return agent.name

    def get_citation_key(self, evidence):
        return evidence.pmid
<commit_msg>Add prefixes for citations and agents<commit_after>import re
import json
import time
import logging
import itertools
from ndex2.nice_cx_network import NiceCXNetwork
from collections import OrderedDict
from indra.statements import *
from indra.databases import context_client, ndex_client, get_identifiers_url, \
    url_prefixes


logger = logging.getLogger(__name__)


class NiceCxAssembler(object):
    def __init__(self, statements=None, name=None):
        self.statements = statements if statements else []
        self.network = NiceCXNetwork()
        self.network.set_network_attribute('name', (name if name else
                                                    'indra_assembled'))
        self.node_keys = {}
        self.citation_keys = {}

    def make_model(self):
        for stmt in self.statements:
            agents = stmt.agent_list()
            not_none_agents = [a for a in agents if a is not None]
            if len(not_none_agents) < 2:
                continue
            for a1, a2 in itertools.combinations(not_none_agents, 2):
                a1_id = self.add_node(a1)
                a2_id = self.add_node(a2)
                edge_id = self.add_edge(a1_id, a2_id, stmt)

        prefixes = {k: v for k, v in url_prefixes.items()}
        prefixes['pubmed'] = 'https://identifiers.org/pubmed'
        self.network.set_network_attribute('@context', json.dumps(prefixes))

    def add_node(self, agent):
        agent_key = self.get_agent_key(agent)
        if agent_key in self.node_keys:
            node_id = self.node_keys[agent_key]
        else:
            node_id = self.network.create_node(agent.name, str(agent.db_refs))
            self.node_keys[agent_key] = node_id
            for db_name, db_id in agent.db_refs.items():
                if db_name in url_prefixes:
                    self.network.add_node_attribute(property_of=node_id,
                                                    name=db_name,
                                                    values='%s:%s' % (db_name,
                                                                      db_id),
                                                    type='string')
        return node_id

    def add_edge(self, a1_id, a2_id, stmt):
        stmt_type = stmt.__class__.__name__
        edge_id = self.network.create_edge(a1_id, a2_id, stmt_type)
        pmids = [ev.pmid for ev in stmt.evidence if ev.pmid is not None]
        self.network.set_edge_attribute(edge_id, 'citation',
                                        ['pubmed:%s' % p for p in pmids],
                                        type='list_of_string')
        return edge_id

    def print_model(self):
        return self.network.to_cx()

    def get_agent_key(self, agent):
        return agent.name

    def get_citation_key(self, evidence):
        return evidence.pmid

{'fraction_non_alphanumeric': 0.08389012620638456, 'fraction_numerical': 0.005567928730512249, 'mean_word_length': 1.857370095440085, 'pattern_counts': {'":': 0, '<': 5, '<?xml version=': 0, '>': 3, 'https://': 1, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 23, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '2341307', 'n_tokens_mistral': 1588, 'n_tokens_neox': 1523, 'n_words': 348}