<commit_before>from django.conf import settings
import simplejson as json
from elasticsearch import Elasticsearch
from maldini import models
from maldini import queues
from maldini import digest as digest_module
from maldini import emails

es = Elasticsearch(settings.ELASTICSEARCH_URL)

def digest(id, verbose):
    try:
        document = models.Document.objects.get(id=id)
    except models.Document.DoesNotExist:
        if verbose: print('MISSING')
        return

    try:
        data = digest_module.digest(document)

    except emails.MissingEmlxPart:
        document.broken = 'missing_emlx_part'
        document.save()
        if verbose: print('missing_emlx_part')
        return

    except emails.PayloadError:
        document.broken = 'payload_error'
        document.save()
        if verbose: print('payload_error')
        return

    except emails.CorruptedFile:
        document.broken = 'corrupted_file'
        document.save()
        if verbose: print('corrupted_file')
        return

    else:
        if document.broken:
            if verbose: print('removing broken flag', document.broken)
            document.broken = ''
            document.save()

    for name, info in data.get('attachments', {}).items():
        child, created = models.Document.objects.update_or_create(
            container=document,
            path=name,
            defaults={
                'disk_size': 0,
                'content_type': info['content_type'],
                'filename': info['filename'],
            },
        )

        if created:
            queues.put('digest', {'id': child.id}, verbose=verbose)
            if verbose: print('new child', child.id)

    models.Digest.objects.update_or_create(
        id=document.id,
        defaults={'data': json.dumps(data)},
    )

    queues.put('index', {'id': document.id}, verbose=verbose)

def index(id, verbose):
    try:
        digest = models.Digest.objects.get(id=id)
    except models.Digest.DoesNotExist:
        if verbose: print('MISSING')
        return

    alldata = json.loads(digest.data)

    data = {
        'title': alldata.get('title'),
        'path': alldata.get('path'),
        'text': alldata.get('text'),
        'subject': alldata.get('subject'),
        'date': alldata.get('date'),
        'people': ' '.join([alldata.get('from', '')] + alldata.get('to', [])),
        'to': alldata.get('to'),
        'from': alldata.get('from'),
        'filetype': alldata.get('type'),
        'sha1': alldata.get('sha1'),
        'md5': alldata.get('md5'),
        'lang': alldata.get('lang'),
        'date_created': alldata.get('date_created'),
        'attachments': bool(alldata.get('attachments')),
        'message_id': alldata.get('message_id'),
        'in_reply_to': alldata.get('in_reply_to'),
        'thread_index': alldata.get('thread_index'),
        'references': alldata.get('references'),
    }

    es.index(
        index=settings.ELASTICSEARCH_INDEX,
        doc_type='doc',
        id=digest.id,
        body=data,
    )<commit_msg>refactor conversion from digest data to index data<commit_after>from django.conf import settings
import simplejson as json
from elasticsearch import Elasticsearch
from maldini import models
from maldini import queues
from maldini import digest as digest_module
from maldini import emails

es = Elasticsearch(settings.ELASTICSEARCH_URL)

def digest(id, verbose):
    try:
        document = models.Document.objects.get(id=id)
    except models.Document.DoesNotExist:
        if verbose: print('MISSING')
        return

    try:
        data = digest_module.digest(document)

    except emails.MissingEmlxPart:
        document.broken = 'missing_emlx_part'
        document.save()
        if verbose: print('missing_emlx_part')
        return

    except emails.PayloadError:
        document.broken = 'payload_error'
        document.save()
        if verbose: print('payload_error')
        return

    except emails.CorruptedFile:
        document.broken = 'corrupted_file'
        document.save()
        if verbose: print('corrupted_file')
        return

    else:
        if document.broken:
            if verbose: print('removing broken flag', document.broken)
            document.broken = ''
            document.save()

    for name, info in data.get('attachments', {}).items():
        child, created = models.Document.objects.update_or_create(
            container=document,
            path=name,
            defaults={
                'disk_size': 0,
                'content_type': info['content_type'],
                'filename': info['filename'],
            },
        )

        if created:
            queues.put('digest', {'id': child.id}, verbose=verbose)
            if verbose: print('new child', child.id)

    models.Digest.objects.update_or_create(
        id=document.id,
        defaults={'data': json.dumps(data)},
    )

    queues.put('index', {'id': document.id}, verbose=verbose)

def index(id, verbose):
    try:
        digest = models.Digest.objects.get(id=id)
    except models.Digest.DoesNotExist:
        if verbose: print('MISSING')
        return

    digest_data = json.loads(digest.data)
    copy_keys = {
        'title',
        'path',
        'text',
        'subject',
        'date',
        'to',
        'from',
        'sha1',
        'md5',
        'lang',
        'date_created',
        'message_id',
        'in_reply_to',
        'thread_index',
        'references',
    }

    data = {key: digest_data.get(key) for key in copy_keys}
    data['filetype'] = digest_data.get('type')
    data['attachments'] = bool(digest_data.get('attachments'))
    data['people'] = ' '.join([digest_data.get('from', '')] + digest_data.get('to', []))

    es.index(
        index=settings.ELASTICSEARCH_INDEX,
        doc_type='doc',
        id=digest.id,
        body=data,
    )
{'fraction_non_alphanumeric': 0.12944983818770225, 'fraction_numerical': 0.0013626298756600238, 'mean_word_length': 2.3515981735159817, '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': '5729620', 'n_tokens_mistral': 1808, 'n_tokens_neox': 1738, 'n_words': 357}