<commit_before>from banal import ensure_list, first
from marshmallow import Schema, post_dump

from aleph.core import es
from aleph.model import Role, Document, Entity, Collection
from aleph.index.core import entities_index, collections_index
from aleph.index.util import unpack_result


class ExpandableSchema(Schema):
    EXPAND = []

    def _get_values(self, obj, field):
        value = obj.get(field)
        if isinstance(value, dict):
            value = value.get('id')

        return [str(v) for v in ensure_list(value)]

    def _type_dispatch(self, type_):
        if type_ in [Collection]:
            return collections_index()
        if type_ in [Document, Entity]:
            return entities_index()
        return type_

    def _resolve_roles(self, cache):
        roles = set()
        for (type_, id_) in cache.keys():
            if type_ == Role:
                roles.add(id_)
        if not len(roles):
            return
        for role in Role.all_by_ids(roles, deleted=True):
            cache[(Role, str(role.id))] = role

    def _resolve_index(self, cache):
        query = []
        for (type_, id_) in cache.keys():
            if type_ == Role:
                continue
            query.append({
                '_index': type_,
                '_doc': 'doc',
                '_id': id_
            })

        if not len(query):
            return

        results = es.mget(body={'docs': query},
                          _source_exclude=['text'])
        for doc in results['docs']:
            cache[(doc['_index'], doc['_id'])] = unpack_result(doc)

    @post_dump(pass_many=True)
    def expand(self, objs, many=False):
        cache = {}
        for obj in ensure_list(objs):
            for (field, type_, _, _, _) in self.EXPAND:
                type_ = self._type_dispatch(type_)
                for key in self._get_values(obj, field):
                    cache[(type_, key)] = None<commit_msg>Remove stray “_doc” mention.<commit_after>from banal import ensure_list, first
from marshmallow import Schema, post_dump

from aleph.core import es
from aleph.model import Role, Document, Entity, Collection
from aleph.index.core import entities_index, collections_index
from aleph.index.util import unpack_result


class ExpandableSchema(Schema):
    EXPAND = []

    def _get_values(self, obj, field):
        value = obj.get(field)
        if isinstance(value, dict):
            value = value.get('id')

        return [str(v) for v in ensure_list(value)]

    def _type_dispatch(self, type_):
        if type_ in [Collection]:
            return collections_index()
        if type_ in [Document, Entity]:
            return entities_index()
        return type_

    def _resolve_roles(self, cache):
        roles = set()
        for (type_, id_) in cache.keys():
            if type_ == Role:
                roles.add(id_)
        if not len(roles):
            return
        for role in Role.all_by_ids(roles, deleted=True):
            cache[(Role, str(role.id))] = role

    def _resolve_index(self, cache):
        query = []
        for (type_, id_) in cache.keys():
            if type_ == Role:
                continue
            query.append({
                '_index': type_,
                '_id': id_
            })

        if not len(query):
            return

        results = es.mget(body={'docs': query},
                          _source_exclude=['text'])
        for doc in results['docs']:
            cache[(doc['_index'], doc['_id'])] = unpack_result(doc)

    @post_dump(pass_many=True)
    def expand(self, objs, many=False):
        cache = {}
        for obj in ensure_list(objs):
            for (field, type_, _, _, _) in self.EXPAND:
                type_ = self._type_dispatch(type_)
                for key in self._get_values(obj, field):
                    cache[(type_, key)] = None
{'fraction_non_alphanumeric': 0.10924586776859505, 'fraction_numerical': 0.0, 'mean_word_length': 1.927437641723356, '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': '5777823', 'n_tokens_mistral': 1129, 'n_tokens_neox': 1090, 'n_words': 308}