<commit_before>import math

from collections import defaultdict
from projex.lazymodule import lazy_import
from projex.locks import ReadWriteLock, ReadLocker, WriteLocker


orb = lazy_import('orb')


class CollectionIterator(object):
    def __init__(self, collection, batch=1):
        self.__collection = collection
        self.__model = collection.model()
        self.__page = 1
        self.__index = -1
        self.__pageSize = batch
        self.__records = []

    def __iter__(self):
        return self

    def next(self):
        self.__index += 1

        # get the next batch of records
        if len(self.__records) in (0, self.__pageSize) and self.__index == len(self.__records):
            sub_collection = self.__collection.page(self.__page, pageSize=self.__pageSize, returning='values')
            self.__records = sub_collection.records()
            self.__page += 1
            self.__index = 0

        # stop the iteration when complete
        if not self.__records or self.__index == len(self.__records):
            raise StopIteration()
        else:
            if self.__collection.context().inflated in (True, None):
                return self.__model.inflate(self.__records[self.__index], context=self.__collection.context())
            else:
                return self.__records[self.__index]


class Collection(object):
    def __json__(self):
        context = self.context()
        expand = context.expandtree(self.__model)

        output = {}

        use_records = False

        if 'count' in expand or context.returning == 'count':
            expand.pop('count', None)
            use_records = True
            output['count'] = self.count()

        if 'ids' in expand or context.returning == 'ids':
            expand.pop('ids', None)
            use_records = True
            output['ids'] = self.ids()

        if 'first' in expand or context.returning == 'first':
            expand.pop('first', None)
            use_records = True
            record = self.first()
            output['first'] = record.__json__() if record else None

        if 'last' in expand or context.returning == 'last':
            expand.pop('last', None)
            use_records = True
            record = self.last()
            output['last'] = record.__json__() if record else None

        if not output or (expand and context.returning not in ('count', 'ids', 'first', 'last')):
            records = [r.__json__() if hasattr(r, '__json__') else r for r in self]
            if not use_records:
                return records
            else:
                output['records'] = records

        return output

    def __init__(self,
                 records=None,
                 model=None,
                 record=None,
                 collector=None,
                 preload=None,
                 **context):
        self.__cacheLock = ReadWriteLock()
        self.__cache = defaultdict(dict)
        self.__preload = preload or {}
        self.__context = orb.Context(**context)
        self.__model = model
        self.__record = record
        self.__collector = collector

        if records is not None and len(records) > 0:
            if self.__model is None:
                self.__model = type(records[0])

            self.__cache['records'][self.__context] = records

    def __len__(self):
        return self.count()

    def __iter__(self):
        if self.isNull():
            return

        context = self.context()

        try:
            with ReadLocker(self.__cacheLock):
                records = self.__cache['records'][context]

        except KeyError:
            try:
                with ReadLocker(self.__cacheLock):
                    raw = self.__preload['records'][context]

            except KeyError:
                conn = context.db.connection()
                raw = conn.select(self.__model, context)

            cache = []
            for x in self._process(raw, context):
                yield x
                cache.append(x)

            with WriteLocker(self.__cacheLock):
                self.__cache['records'][context] = cache

        else:
            for r in records:
                yield r

    def __getitem__(self, index):
        if isinstance(index, slice):
            return self.copy(start=index.start, limit=(index.stop - index.start))
        else:
            record = self.at(index)
            if not record:
                raise IndexError(index)
            else:
                return record

    def _process(self, raw, context):
        if context.inflated in (True, None) and context.returning != 'values':
            for x in raw or []:
                yield self.__model.inflate(x, context=context)

        elif context.columns:
            schema = self.__model.schema()
            if context.returning == 'values':
                if len(context.columns) == 1:
                    col = schema.column(context.columns[0])
                    for x in raw or []:
                        yield x[col.field()]

                else:
                    cols = [schema.column(col) for col in context.columns]
                    for x in raw or []:
                        yield tuple(x.get(col.field()) for col in cols)
            else:
                cols = [schema.column(col) for col in context.columns]
                for x in raw or []:
                    yield {col.field(): x.get(col.field()) for col in cols}
        else:<commit_msg>* added in support for returning data objects<commit_after>import math

from collections import defaultdict
from projex.lazymodule import lazy_import
from projex.locks import ReadWriteLock, ReadLocker, WriteLocker


orb = lazy_import('orb')


class CollectionIterator(object):
    def __init__(self, collection, batch=1):
        self.__collection = collection
        self.__model = collection.model()
        self.__page = 1
        self.__index = -1
        self.__pageSize = batch
        self.__records = []

    def __iter__(self):
        return self

    def next(self):
        self.__index += 1

        # get the next batch of records
        if len(self.__records) in (0, self.__pageSize) and self.__index == len(self.__records):
            sub_collection = self.__collection.page(self.__page, pageSize=self.__pageSize, returning='values')
            self.__records = sub_collection.records()
            self.__page += 1
            self.__index = 0

        # stop the iteration when complete
        if not self.__records or self.__index == len(self.__records):
            raise StopIteration()
        else:
            if self.__collection.context().inflated in (True, None):
                return self.__model.inflate(self.__records[self.__index], context=self.__collection.context())
            else:
                return self.__records[self.__index]


class Collection(object):
    def __json__(self):
        context = self.context()
        expand = context.expandtree(self.__model)

        output = {}

        use_records = False

        if 'count' in expand or context.returning == 'count':
            expand.pop('count', None)
            use_records = True
            output['count'] = self.count()

        if 'ids' in expand or context.returning == 'ids':
            expand.pop('ids', None)
            use_records = True
            output['ids'] = self.ids()

        if 'first' in expand or context.returning == 'first':
            expand.pop('first', None)
            use_records = True
            record = self.first()
            output['first'] = record.__json__() if record else None

        if 'last' in expand or context.returning == 'last':
            expand.pop('last', None)
            use_records = True
            record = self.last()
            output['last'] = record.__json__() if record else None

        if not output or (expand and context.returning not in ('count', 'ids', 'first', 'last')):
            records = [r.__json__() if hasattr(r, '__json__') else r for r in self]
            if not use_records:
                return records
            else:
                output['records'] = records

        return output

    def __init__(self,
                 records=None,
                 model=None,
                 record=None,
                 collector=None,
                 preload=None,
                 **context):
        self.__cacheLock = ReadWriteLock()
        self.__cache = defaultdict(dict)
        self.__preload = preload or {}
        self.__context = orb.Context(**context)
        self.__model = model
        self.__record = record
        self.__collector = collector

        if records is not None and len(records) > 0:
            if self.__model is None:
                self.__model = type(records[0])

            self.__cache['records'][self.__context] = records

    def __len__(self):
        return self.count()

    def __iter__(self):
        if self.isNull():
            return

        context = self.context()

        try:
            with ReadLocker(self.__cacheLock):
                records = self.__cache['records'][context]

        except KeyError:
            try:
                with ReadLocker(self.__cacheLock):
                    raw = self.__preload['records'][context]

            except KeyError:
                conn = context.db.connection()
                raw = conn.select(self.__model, context)

            cache = []
            for x in self._process(raw, context):
                yield x
                cache.append(x)

            with WriteLocker(self.__cacheLock):
                self.__cache['records'][context] = cache

        else:
            for r in records:
                yield r

    def __getitem__(self, index):
        if isinstance(index, slice):
            return self.copy(start=index.start, limit=(index.stop - index.start))
        else:
            record = self.at(index)
            if not record:
                raise IndexError(index)
            else:
                return record

    def _process(self, raw, context):
        if context.inflated in (True, None) and context.returning not in ('values', 'data'):
            for x in raw or []:
                yield self.__model.inflate(x, context=context)

        elif context.columns:
            schema = self.__model.schema()
            if context.returning == 'values':
                if len(context.columns) == 1:
                    col = schema.column(context.columns[0])
                    for x in raw or []:
                        yield x[col.field()]

                else:
                    cols = [schema.column(col) for col in context.columns]
                    for x in raw or []:
                        yield tuple(x.get(col.field()) for col in cols)
            else:
                cols = [schema.column(col) for col in context.columns]
                for x in raw or []:
                    yield {col.field(): x.get(col.field()) for col in cols}
        else:
{'fraction_non_alphanumeric': 0.1008907471368842, 'fraction_numerical': 0.0019996364297400473, 'mean_word_length': 1.6538832609744332, 'pattern_counts': {'":': 0, '<': 3, '<?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': '3181994', 'n_tokens_mistral': 2880, 'n_tokens_neox': 2749, 'n_words': 785}