<commit_before>import random

import openpathsampling as paths
from openpathsampling.todict import OPSNamed

import logging
logger = logging.getLogger(__name__)

class SampleKeyError(Exception):
    def __init__(self, key, sample, sample_key):
        self.key = key
        self.sample = sample
        self.sample_key = sample_key
        self.msg = (str(self.key) + " does not match " + str(self.sample_key)
                    + " from " + str(self.sample))

class SampleSet(OPSNamed):
    '''
    SampleSet is essentially a list of samples, with a few conveniences.  It
    can be treated as a list of samples (using, e.g., .append), or as a
    dictionary of ensembles mapping to a list of samples, or as a dictionary
    of replica IDs to samples. Any type is allowed as a replica ID except
    Sample or Ensemble.

    The dictionaries ensemble_dict and replica_dict are conveniences which
    should be kept consistent by any method which modifies the container.
    They do not need to be stored.

    Note
    ----
        Current implementation is as an unordered set. Therefore we don't
        have some of the convenient tools in Python sequences (e.g.,
        slices). On the other hand, I'm not sure whether that is meaningful
        here.
        Since replicas are integers we add slicing/ranges for replicas. In addition
        we support any iterable as input in __getitem__ an it will return an iterable
        over the results. This makes it possible to write `sset[0:5]` to get a list
        of of ordered samples by replica_id, or sset[list_of_ensembles]


    Attributes
    ----------
    samples : list of Sample
        The samples included in this set.
    ensemble_dict : dict
        A dictionary with Ensemble objects as keys and lists of Samples as
        values.
    replica_dict : dict
        A dictionary with replica IDs as keys and lists of Samples as values
    '''

    def __init__(self, samples, movepath=None):
        super(SampleSet, self).__init__()
        self.samples = []
        self.ensemble_dict = {}
        self.replica_dict = {}
        self.extend(samples)
        if movepath is None:
            self.movepath = paths.EmptyPathMoveChange()
        else:
            self.movepath = movepath

    def __getitem__(self, key):
        if hasattr(key, '__iter__'):
            return (self[element] for element in key)

        elif isinstance(key, paths.Ensemble):
            return random.choice(self.ensemble_dict[key])
        else:
            return random.choice(self.replica_dict[key])

    def __setitem__(self, key, value):
        # first, we check whether the key matches the sample: if no, KeyError
        if isinstance(key, paths.Ensemble):
            if key != value.ensemble:
                raise SampleKeyError(key, value, value.ensemble)
        else:
            if key != value.replica:
                raise SampleKeyError(key, value, value.replica)

        if value in self.samples:
            # if value is already in this, we don't need to do anything
            return
        # Setting works by replacing one with the same key. We pick one with
        # this key at random (using __getitem__), delete it, and then append
        # the new guy. If nothing exists with the desired key, this is the
        # same as append.
        try:
            dead_to_me = self[key]
        except KeyError:
            dead_to_me = None
        if dead_to_me is not None:<commit_msg>slice support for samplesets<commit_after>import random

import openpathsampling as paths
from openpathsampling.todict import OPSNamed

import logging
logger = logging.getLogger(__name__)

class SampleKeyError(Exception):
    def __init__(self, key, sample, sample_key):
        self.key = key
        self.sample = sample
        self.sample_key = sample_key
        self.msg = (str(self.key) + " does not match " + str(self.sample_key)
                    + " from " + str(self.sample))

class SampleSet(OPSNamed):
    '''
    SampleSet is essentially a list of samples, with a few conveniences.  It
    can be treated as a list of samples (using, e.g., .append), or as a
    dictionary of ensembles mapping to a list of samples, or as a dictionary
    of replica IDs to samples. Any type is allowed as a replica ID except
    Sample or Ensemble.

    The dictionaries ensemble_dict and replica_dict are conveniences which
    should be kept consistent by any method which modifies the container.
    They do not need to be stored.

    Note
    ----
        Current implementation is as an unordered set. Therefore we don't
        have some of the convenient tools in Python sequences (e.g.,
        slices). On the other hand, I'm not sure whether that is meaningful
        here.
        Since replicas are integers we add slicing/ranges for replicas. In addition
        we support any iterable as input in __getitem__ an it will return an iterable
        over the results. This makes it possible to write `sset[0:5]` to get a list
        of of ordered samples by replica_id, or sset[list_of_ensembles].
        replica_ids can be any number do not have to be subsequent to slicing does not
        make sense and we ignore it. We will also ignore missing replica_ids. A slice
        `1:5` will return all existing replica ids >=1 and <5. If you want exactly
        all replicas from 1 to 4 use `sset[xrange(1,5)]`


    Attributes
    ----------
    samples : list of Sample
        The samples included in this set.
    ensemble_dict : dict
        A dictionary with Ensemble objects as keys and lists of Samples as
        values.
    replica_dict : dict
        A dictionary with replica IDs as keys and lists of Samples as values
    '''

    def __init__(self, samples, movepath=None):
        super(SampleSet, self).__init__()
        self.samples = []
        self.ensemble_dict = {}
        self.replica_dict = {}
        self.extend(samples)
        if movepath is None:
            self.movepath = paths.EmptyPathMoveChange()
        else:
            self.movepath = movepath

    def __getitem__(self, key):
        if hasattr(key, '__iter__'):
            return (self[element] for element in key)
        elif type(key) is slice:
            rep_idxs = filter(
                lambda x :
                    (key.start is None or x >= key.start) and
                    (key.stop is None or x < key.stop),
                sorted(self.replica_dict.keys())

            )

            return (self[element] for element in rep_idxs)
        elif isinstance(key, paths.Ensemble):
            return random.choice(self.ensemble_dict[key])
        elif type(key) is int:
            return random.choice(self.replica_dict[key])

    def __setitem__(self, key, value):
        # first, we check whether the key matches the sample: if no, KeyError
        if isinstance(key, paths.Ensemble):
            if key != value.ensemble:
                raise SampleKeyError(key, value, value.ensemble)
        else:
            if key != value.replica:
                raise SampleKeyError(key, value, value.replica)

        if value in self.samples:
            # if value is already in this, we don't need to do anything
            return
        # Setting works by replacing one with the same key. We pick one with
        # this key at random (using __getitem__), delete it, and then append
        # the new guy. If nothing exists with the desired key, this is the
        # same as append.
        try:
            dead_to_me = self[key]
        except KeyError:
            dead_to_me = None
        if dead_to_me is not None:
{'fraction_non_alphanumeric': 0.0677698975571316, 'fraction_numerical': 0.0015760441292356187, 'mean_word_length': 2.210370994940978, 'pattern_counts': {'":': 0, '<': 5, '<?xml version=': 0, '>': 5, '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': '893543', 'n_tokens_mistral': 2090, 'n_tokens_neox': 2032, 'n_words': 850}