<commit_before>import uuid
import contextlib

from tvtk.api import tvtk

from simphony.cuds.abstractparticles import ABCParticles
from simphony.cuds.particles import Particle, Bond
from simphony.core.data_container import DataContainer
from simphony_mayavi.core.api import (
    CubaData, CellCollection, supported_cuba, mergedocs)


VTK_POLY_LINE = 4


@mergedocs(ABCParticles)
class VTKParticles(ABCParticles):

    def __init__(self, name, data=None, data_set=None):
        self.name = name
        self._data = DataContainer() if data is None else DataContainer(data)
        #: The mapping from uid to point index
        self.particle2index = {}
        #: The reverse mapping from index to point uid
        self.index2particle = {}
        #: The mapping from uid to bond index
        self.bond2index = {}
        #: The reverse mapping from index to bond uid
        self.index2bond = {}

        # Setup the data_set
        if data_set is None:
            points = tvtk.Points()
            # Need to initialise lines with empty so that we
            # do not get the shared CellArray
            data_set = tvtk.PolyData(points=points, lines=[])
            # To use get_cell we need to have at least one point in the
            # dataset.
            points.append((0, 0, 0))
            self.initialized = True
        else:
            self.initialized = False
            for index in xrange(data_set.number_of_points):
                uid = uuid.uuid4()
                self.particle2index[uid] = index
                self.index2particle[index] = uid
            for index in xrange(data_set.number_of_cells):
                uid = uuid.uuid4()
                self.bond2index[uid] = index
                self.index2bond[index] = uid

        #: The vtk.PolyData dataset
        self.data_set = data_set

        #: The currently supported and stored CUBA keywords.
        self.supported_cuba = supported_cuba()
        #: Easy access to the vtk PointData structure
        self.point_data = CubaData(
            data_set.point_data, stored_cuba=self.supported_cuba)
        #: Easy access to the vtk CellData structure
        self.bond_data = CubaData(
            data_set.cell_data, stored_cuba=self.supported_cuba)
        #: Easy access to the lines vtk CellArray structure
        self.bonds = CellCollection(data_set.lines)

    @property
    def data(self):
        return DataContainer(self._data)

    @data.setter
    def data(self, value):
        self._data = DataContainer(value)

    # Particle operations ####################################################

    def add_particle(self, particle):
        data_set = self.data_set
        points = data_set.points
        particle2index = self.particle2index
        with self._add_item(particle, particle2index) as item:
            if self.initialized:
                # We remove the dummy point
                self.data_set.points = tvtk.Points()
                points = data_set.points
                self.initialized = False
            index = points.insert_next_point(item.coordinates)
            particle2index[item.uid] = index
            self.index2particle[index] = item.uid
            self.point_data.append(item.data)
        return item.uid

    def get_particle(self, uid):
        index = int(self.particle2index[uid])
        return Particle(
            uid=uid,
            coordinates=self.data_set.points[index],
            data=self.point_data[index])
<commit_msg>add docstring<commit_after>import uuid
import contextlib

from tvtk.api import tvtk

from simphony.cuds.abstractparticles import ABCParticles
from simphony.cuds.particles import Particle, Bond
from simphony.core.data_container import DataContainer
from simphony_mayavi.core.api import (
    CubaData, CellCollection, supported_cuba, mergedocs)


VTK_POLY_LINE = 4


@mergedocs(ABCParticles)
class VTKParticles(ABCParticles):

    def __init__(self, name, data=None, data_set=None):
        self.name = name
        self._data = DataContainer() if data is None else DataContainer(data)
        #: The mapping from uid to point index
        self.particle2index = {}
        #: The reverse mapping from index to point uid
        self.index2particle = {}
        #: The mapping from uid to bond index
        self.bond2index = {}
        #: The reverse mapping from index to bond uid
        self.index2bond = {}

        # Setup the data_set
        if data_set is None:
            points = tvtk.Points()
            # Need to initialise lines with empty so that we
            # do not get the shared CellArray
            data_set = tvtk.PolyData(points=points, lines=[])
            # To use get_cell we need to have at least one point in the
            # dataset.
            points.append((0, 0, 0))
            self.initialized = True
        else:
            self.initialized = False
            for index in xrange(data_set.number_of_points):
                uid = uuid.uuid4()
                self.particle2index[uid] = index
                self.index2particle[index] = uid
            for index in xrange(data_set.number_of_cells):
                uid = uuid.uuid4()
                self.bond2index[uid] = index
                self.index2bond[index] = uid

        #: The vtk.PolyData dataset
        self.data_set = data_set

        #: The currently supported and stored CUBA keywords.
        self.supported_cuba = supported_cuba()
        #: Easy access to the vtk PointData structure
        self.point_data = CubaData(
            data_set.point_data, stored_cuba=self.supported_cuba)
        #: Easy access to the vtk CellData structure
        self.bond_data = CubaData(
            data_set.cell_data, stored_cuba=self.supported_cuba)
        #: Easy access to the lines vtk CellArray structure
        self.bonds = CellCollection(data_set.lines)

    @property
    def data(self):
        """ The container data
        """
        return DataContainer(self._data)

    @data.setter
    def data(self, value):
        self._data = DataContainer(value)

    # Particle operations ####################################################

    def add_particle(self, particle):
        data_set = self.data_set
        points = data_set.points
        particle2index = self.particle2index
        with self._add_item(particle, particle2index) as item:
            if self.initialized:
                # We remove the dummy point
                self.data_set.points = tvtk.Points()
                points = data_set.points
                self.initialized = False
            index = points.insert_next_point(item.coordinates)
            particle2index[item.uid] = index
            self.index2particle[index] = item.uid
            self.point_data.append(item.data)
        return item.uid

    def get_particle(self, uid):
        index = int(self.particle2index[uid])
        return Particle(
            uid=uid,
            coordinates=self.data_set.points[index],
            data=self.point_data[index])

{'fraction_non_alphanumeric': 0.08923734853884534, 'fraction_numerical': 0.005702066999287242, 'mean_word_length': 2.164636896707262, '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': '6944329', 'n_tokens_mistral': 1996, 'n_tokens_neox': 1895, 'n_words': 516}