<commit_before>from pathlib import Path
from io import StringIO
import json
from django.db import models, transaction
from django.contrib.postgres.fields import JSONField
from django.conf import settings

def cache(model, keyfunc):

    def decorator(func):

        if not settings.MALDINI_CACHE:
            return func

        @transaction.atomic
        def wrapper(*args, **kwargs):
            key = keyfunc(*args, **kwargs)

            row, created = model.objects.get_or_create(pk=key)
            if not created:
                return json.loads(row.value)

            value = func(*args, **kwargs)

            row.value = json.dumps(value)
            row.save()

            return value

        return wrapper

    return decorator

class EmailCache(models.Model):
    id = models.IntegerField(primary_key=True)
    value = models.TextField()
    time = models.DateTimeField(auto_now=True)

from . import emails  # moved here because circular import

class Document(models.Model):
    container = models.ForeignKey('Document', null=True)
    path = models.CharField(max_length=4000)
    content_type = models.CharField(max_length=100, blank=True)
    filename = models.CharField(max_length=1000)
    disk_size = models.BigIntegerField()
    md5 = models.CharField(max_length=40, blank=True, db_index=True)
    sha1 = models.CharField(max_length=50, blank=True, db_index=True)
    broken = models.CharField(max_length=100, blank=True)
    rev = models.IntegerField(null=True)
    flags = JSONField(default=dict, blank=True)

    class Meta:
        # TODO: constraint does not apply to container=None rows
        unique_together = ('container', 'path')

    @property
    def absolute_path(self):
        return Path(settings.MALDINI_ROOT) / self.path

    def open(self):
        if self.content_type == 'application/x-directory':
            return StringIO()

        if self.container is None:
            return self.absolute_path.open('rb')

        else:
            if emails.is_email(self.container):
                return emails.get_email_part(self.container, self.path)

        raise RuntimeError

class Ocr(models.Model):
    tag = models.CharField(max_length=100)
    md5 = models.CharField(max_length=40, db_index=True)
    path = models.CharField(max_length=4000)
    text = models.TextField(blank=True)

    class Meta:
        unique_together = ('tag', 'md5')

    @property
    def absolute_path(self):
        return Path(settings.MALDINI_OCR_ROOT) / self.tag / self.path

class Digest(models.Model):
    id = models.IntegerField(primary_key=True)
    data = models.TextField()
    updated_at = models.DateTimeField(auto_now=True)

class FolderMark(models.Model):
    path = models.CharField(max_length=4000, unique=True, db_index=True)

class Job(models.Model):
    queue = models.CharField(max_length=100)
    data = JSONField(null=True)
    started = models.BooleanField(default=False)

    class Meta:
        unique_together = ('queue', 'data')
        index_together = ('queue', 'started')

class TikaCache(models.Model):
    sha1 = models.CharField(max_length=50, primary_key=True)
    value = models.TextField()
    time = models.DateTimeField(auto_now=True)

class TikaLangCache(models.Model):
    sha1 = models.CharField(max_length=50, primary_key=True)
    value = models.CharField(max_length=20)
    time = models.DateTimeField(auto_now=True)
<commit_msg>use BytesIO instead of StringIO<commit_after>from pathlib import Path
from io import BytesIO
import json
from django.db import models, transaction
from django.contrib.postgres.fields import JSONField
from django.conf import settings

def cache(model, keyfunc):

    def decorator(func):

        if not settings.MALDINI_CACHE:
            return func

        @transaction.atomic
        def wrapper(*args, **kwargs):
            key = keyfunc(*args, **kwargs)

            row, created = model.objects.get_or_create(pk=key)
            if not created:
                return json.loads(row.value)

            value = func(*args, **kwargs)

            row.value = json.dumps(value)
            row.save()

            return value

        return wrapper

    return decorator

class EmailCache(models.Model):
    id = models.IntegerField(primary_key=True)
    value = models.TextField()
    time = models.DateTimeField(auto_now=True)

from . import emails  # moved here because circular import

class Document(models.Model):
    container = models.ForeignKey('Document', null=True)
    path = models.CharField(max_length=4000)
    content_type = models.CharField(max_length=100, blank=True)
    filename = models.CharField(max_length=1000)
    disk_size = models.BigIntegerField()
    md5 = models.CharField(max_length=40, blank=True, db_index=True)
    sha1 = models.CharField(max_length=50, blank=True, db_index=True)
    broken = models.CharField(max_length=100, blank=True)
    rev = models.IntegerField(null=True)
    flags = JSONField(default=dict, blank=True)

    class Meta:
        # TODO: constraint does not apply to container=None rows
        unique_together = ('container', 'path')

    @property
    def absolute_path(self):
        return Path(settings.MALDINI_ROOT) / self.path

    def open(self):
        if self.content_type == 'application/x-directory':
            return BytesIO()

        if self.container is None:
            return self.absolute_path.open('rb')

        else:
            if emails.is_email(self.container):
                return emails.get_email_part(self.container, self.path)

        raise RuntimeError

class Ocr(models.Model):
    tag = models.CharField(max_length=100)
    md5 = models.CharField(max_length=40, db_index=True)
    path = models.CharField(max_length=4000)
    text = models.TextField(blank=True)

    class Meta:
        unique_together = ('tag', 'md5')

    @property
    def absolute_path(self):
        return Path(settings.MALDINI_OCR_ROOT) / self.tag / self.path

class Digest(models.Model):
    id = models.IntegerField(primary_key=True)
    data = models.TextField()
    updated_at = models.DateTimeField(auto_now=True)

class FolderMark(models.Model):
    path = models.CharField(max_length=4000, unique=True, db_index=True)

class Job(models.Model):
    queue = models.CharField(max_length=100)
    data = JSONField(null=True)
    started = models.BooleanField(default=False)

    class Meta:
        unique_together = ('queue', 'data')
        index_together = ('queue', 'started')

class TikaCache(models.Model):
    sha1 = models.CharField(max_length=50, primary_key=True)
    value = models.TextField()
    time = models.DateTimeField(auto_now=True)

class TikaLangCache(models.Model):
    sha1 = models.CharField(max_length=50, primary_key=True)
    value = models.CharField(max_length=20)
    time = models.DateTimeField(auto_now=True)

{'fraction_non_alphanumeric': 0.10456942003514938, 'fraction_numerical': 0.013473930872876391, 'mean_word_length': 3.2284829721362227, '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': '5477776', 'n_tokens_mistral': 2136, 'n_tokens_neox': 2035, 'n_words': 438}