<commit_before>"""Manage collections of Connections.
"""

import collections
import numpy as np

from nengo.utils.builder import full_transform


TransformFunctionPair = collections.namedtuple(
    'TransformFunctionPair', ['transform', 'function'])


TransformFunctionWithSolverEvalPoints = collections.namedtuple(
    'TransformFunctionWithSolverEvalPoints',
    ['transform', 'function', 'solver', 'eval_points'])


class Connections(object):
    def __init__(self, connections=[]):
        self._connection_indices = dict()
        self._source = None
        self.transforms_functions = list()

        for connection in connections:
            self.add_connection(connection)

    def add_connection(self, connection):
        # Ensure that this Connection collection is only for connections from
        # the same source object
        if self._source is None:
            self._source = connection.pre
        assert(self._source == connection.pre)

        # Get the index of the connection if the same transform and function
        # have already been added, otherwise add the transform/function pair
        transform = full_transform(connection, allow_scalars=False)
        connection_entry = self._make_connection_entry(connection, transform)
        for (i, tf) in enumerate(self.transforms_functions):
            if self._are_compatible_connections(tf, connection_entry):
                index = i
                break
        else:
            self.transforms_functions.append(connection_entry)
            index = len(self.transforms_functions) - 1

        self._connection_indices[connection] = index

    def contains_compatible_connection(self, connection):
        if self._source is None or self._source != connection.pre:
            return False

        for tf in self.transforms_functions:
            if self._are_compatible_connections(tf, connection):
                return True
        return False

    def _are_compatible_connections(self, c1, c2):
        return (np.all(c1.transform == c2.transform) and
                c1.function == c2.function)

    def _make_connection_entry(self, connection, transform):
        return TransformFunctionPair(transform, connection.function)

    @property
    def width(self):
        return sum([t.transform.shape[0] for t in self.transforms_functions])

    def get_connection_offset(self, connection):
        i = self[connection]
        return sum([t.transform.shape[0] for t in
                    self.transforms_functions[:i]])

    def __len__(self):
        return len(self.transforms_functions)

    def __iter__(self):
        return iter(self._connection_indices)

    def __getitem__(self, connection):
        return self._connection_indices[connection]


class ConnectionsWithSolvers(Connections):
    def _are_compatible_connections(self, c1, c2):
        return (np.all(c1.transform == c2.transform) and
                np.all(c1.eval_points == c2.eval_points) and
                c1.solver == c2.solver and
                c1.function == c2.function)

    def _make_connection_entry(self, connection, transform):
        return TransformFunctionWithSolverEvalPoints(
            transform, connection.function, connection.solver,
            connection.eval_points)


class ConnectionBank(object):
    """Represents a set of Connection collections which need not necessarily
    share the same source object.
    """

    def __init__(self, connections, collector_type=Connections):
        self._connections = collections.defaultdict(collector_type)

        for connection in connections:
            self.add_connection(connection)

    def add_connection(self, connection):
        self._connections[connection.pre].add_connection(connection)

    def get_connection_offset(self, connection):
        offset = 0
        for connections in self._connections.values():
            if connection in connections:<commit_msg>Connection utilities accept keyspaces.<commit_after>"""Manage collections of Connections.
"""

import collections
import numpy as np

from nengo.utils.builder import full_transform

from . import keyspaces
default_keyspace = keyspaces.nengo_default()


TransformFunctionKeyspace = collections.namedtuple(
    'TransformFunctionPair', ['transform', 'function', 'keyspace'])


TransformFunctionWithSolverEvalPoints = collections.namedtuple(
    'TransformFunctionWithSolverEvalPoints',
    ['transform', 'function', 'solver', 'eval_points', 'keyspace'])


class Connections(object):
    def __init__(self, connections=[]):
        self._connection_indices = dict()
        self._source = None
        self.transforms_functions = list()

        for connection in connections:
            self.add_connection(connection)

    def add_connection(self, connection, keyspace=default_keyspace):
        # Ensure that this Connection collection is only for connections from
        # the same source object
        if self._source is None:
            self._source = connection.pre
        assert(self._source == connection.pre)

        # Get the index of the connection if the same transform and function
        # have already been added, otherwise add the transform/function pair
        transform = full_transform(connection, allow_scalars=False)
        connection_entry = self._make_connection_entry(
            connection, transform, keyspace)

        for (i, tf) in enumerate(self.transforms_functions):
            if self._are_compatible_connections(tf, connection_entry):
                index = i
                break
        else:
            self.transforms_functions.append(connection_entry)
            index = len(self.transforms_functions) - 1

        self._connection_indices[connection] = index

    def contains_compatible_connection(self, connection,
                                       keyspace=default_keyspace):
        if self._source is None or self._source != connection.pre:
            return False

        for tf in self.transforms_functions:
            if self._are_compatible_connections(tf, connection):
                return True
        return False

    def _are_compatible_connections(self, c1, c2):
        return (np.all(c1.transform == c2.transform) and
                c1.function == c2.function)

    def _make_connection_entry(self, connection, transform,
                               keyspace=default_keyspace):
        return TransformFunctionKeyspace(transform, connection.function,
                                         keyspace)

    @property
    def width(self):
        return sum([t.transform.shape[0] for t in self.transforms_functions])

    def get_connection_offset(self, connection):
        i = self[connection]
        return sum([t.transform.shape[0] for t in
                    self.transforms_functions[:i]])

    def __len__(self):
        return len(self.transforms_functions)

    def __iter__(self):
        return iter(self._connection_indices)

    def __getitem__(self, connection):
        return self._connection_indices[connection]


class ConnectionsWithSolvers(Connections):
    def _are_compatible_connections(self, c1, c2):
        return (np.all(c1.transform == c2.transform) and
                np.all(c1.eval_points == c2.eval_points) and
                c1.solver == c2.solver and
                c1.function == c2.function)

    def _make_connection_entry(self, connection, transform,
                               keyspace=default_keyspace):
        return TransformFunctionWithSolverEvalPoints(
            transform, connection.function, connection.solver,
            connection.eval_points, keyspace)


class ConnectionBank(object):
    """Represents a set of Connection collections which need not necessarily
    share the same source object.
    """

    def __init__(self, connections, collector_type=Connections):
        self._connections = collections.defaultdict(collector_type)

        for connection in connections:
            self.add_connection(connection)

    def add_connection(self, connection):
        self._connections[connection.pre].add_connection(connection)

    def get_connection_offset(self, connection):
        offset = 0
        for connections in self._connections.values():
            if connection in connections:
{'fraction_non_alphanumeric': 0.07898875045361074, 'fraction_numerical': 0.00483851457602516, 'mean_word_length': 2.803127874885005, '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': 3, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '2747451', 'n_tokens_mistral': 2119, 'n_tokens_neox': 2026, 'n_words': 533}