<commit_before>import elasticsearch.helpers as eh

from esengine.bases.document import BaseDocument
from esengine.bases.metaclass import ModelMetaclass
from esengine.utils import validate_client


class Document(BaseDocument):
    __metaclass__ = ModelMetaclass
    _autoid = True

    @classmethod
    def get_es(cls, es):
        """
        This proxy-method allows the client overwrite
        and the use of a default client for a document.
        Document transport methods should use cls.get_es(es).method()
        This method also validades that the connection is a valid ES client.
        :return: elasticsearch.ElasticSearch() instance or equivalent client
        """
        validate_client(es)
        return es

    def save(self, es=None):
        doc = self.to_dict()
        saved_document = self.get_es(es).index(
            index=self._index,
            doc_type=self._doctype,
            id=self.id,
            body=doc
        )
        if saved_document.get('created'):
            self.id = saved_document['_id']

    @classmethod
    def get(cls, id=None, ids=None, es=None):
        es = cls.get_es(es)
        if id is not None and ids is not None:
            raise ValueError('id and ids can not be passed together.')
        if id is not None:
            res = es.get(index=cls._index,
                         doc_type=cls._doctype,
                         id=id)
            return cls.from_dict(dct=res['_source'])
        if ids is not None:
            query = {
                "query": {
                    "filtered": {
                        "query": {"match_all": {}},
                        "filter": {
                            "ids": {
                                "values": list(ids)
                            }
                        }
                    }
                }}
            resp = es.search(
                index=cls._index,
                doc_type=cls._doctype,
                body=query,
                size=len(ids)
            )
            result = []
            for obj in resp['hits']['hits']:
                result.append(cls.from_dict(dct=obj['_source']['doc']))
            return result
<commit_msg>get returns only one doc and filter returns a list<commit_after>import elasticsearch.helpers as eh

from esengine.bases.document import BaseDocument
from esengine.bases.metaclass import ModelMetaclass
from esengine.utils import validate_client


class Document(BaseDocument):
    __metaclass__ = ModelMetaclass
    _autoid = True

    @classmethod
    def get_es(cls, es):
        """
        This proxy-method allows the client overwrite
        and the use of a default client for a document.
        Document transport methods should use cls.get_es(es).method()
        This method also validades that the connection is a valid ES client.
        :return: elasticsearch.ElasticSearch() instance or equivalent client
        """
        validate_client(es)
        return es

    def save(self, es=None):
        doc = self.to_dict()
        saved_document = self.get_es(es).index(
            index=self._index,
            doc_type=self._doctype,
            id=self.id,
            body=doc
        )
        if saved_document.get('created'):
            self.id = saved_document['_id']

    @classmethod
    def get(cls, id, es=None, **kwargs):
        es = cls.get_es(es)
        res = es.get(index=cls._index,
                     doc_type=cls._doctype,
                     id=id,
                     **kwargs)
        return cls.from_dict(dct=res['_source'])

    @classmethod
    def filter(cls, es=None, **kwargs):
        es = cls.get_es(es)
        query = {
            "query": {
                "filtered": {
                    "query": {"match_all": {}},
                    "filter": {
                        "ids": {
                            "values": None
                        }
                    }
                }
            }}
        resp = es.search(
            index=cls._index,
            doc_type=cls._doctype,
            body=query,
            size=len(ids)
        )
        result = []
        for obj in resp['hits']['hits']:
            result.append(cls.from_dict(dct=obj['_source']['doc']))
        return result

{'fraction_non_alphanumeric': 0.09575217085191269, 'fraction_numerical': 0.0, 'mean_word_length': 1.5011737089201878, 'pattern_counts': {'":': 14, '<': 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': '1507649', 'n_tokens_mistral': 1172, 'n_tokens_neox': 1135, 'n_words': 288}