<commit_before>import gzip
import json
from django.db import models
from django.db import transaction


BULK_CREATE_BATCH_SIZE = 10000


class Annotation(object):
    """
    Can contain multiple Chunk_Feature objects merged together.
    """

    def __init__(self, base_first, base_last, chunk_feature, fragment=None):
        self.base_first = base_first
        self.base_last = base_last
        self.feature = chunk_feature.feature
        self.feature_base_first = chunk_feature.feature_base_first
        self.feature_base_last = chunk_feature.feature_base_last
        self.fragment = fragment

    @property
    def feature_name(self):
        if (
            self.feature_base_first != 1
            or self.feature_base_last != self.feature.length
        ):
            return "%s[%s:%s]" % (
                self.feature.name,
                self.feature_base_first,
                self.feature_base_last,
            )
        else:
            return self.feature.name

    def __str__(self):
        s = []
        s.append(self.feature_name)
        s.append(self.feature.type)
        if self.feature.strand == 1:
            s.append("+")
        else:
            s.append("-")
        return ", ".join(s)

    @staticmethod
    def from_chunk_feature_and_location_array(chunk_feature_locs):
        """
        chunk_feature_locs: an array of (Chunk_Feature, Fragment_Chunk_Location) tuples.
        """

        chunk_feature_locs = sorted(
            chunk_feature_locs,
            key=lambda t: (t[0].feature.id, t[1].base_first,
                           (t[0].feature.strand if t[0].feature.strand is not None else 1)
                           * t[0].feature_base_first)
        )

        annotations = []
        for cf, fcl in chunk_feature_locs:
            if (
                len(annotations) > 0
                and annotations[-1].feature.id == cf.feature_id
                and (((annotations[-1].feature.strand is None or annotations[-1].feature.strand > 0)
                      and annotations[-1].feature_base_last == cf.feature_base_first - 1)
                     or ((annotations[-1].feature.strand is not None
                          and annotations[-1].feature.strand < 0)
                         and annotations[-1].feature_base_first == cf.feature_base_last + 1))
                and annotations[-1].base_last == fcl.base_first - 1
            ):
                # merge annotation
                annotations[-1].base_last = fcl.base_last
                if (annotations[-1].feature.strand is None or annotations[-1].feature.strand > 0):
                    annotations[-1].feature_base_last = cf.feature_base_last
                else:
                    annotations[-1].feature_base_first = cf.feature_base_first
            else:
                annotations.append(
                    Annotation(
                        base_first=fcl.base_first,
                        base_last=fcl.base_last,
                        chunk_feature=cf,
                        fragment=fcl.fragment,
                    )
                )

        return sorted(annotations, key=lambda a: a.base_first)


class ChunkReference(object):
    """
    Parent class for storing chunk sequences in reference files
    1-based indexing example: (2, 4) of 5'-A(TCG)GCTA-3'
    """

    @staticmethod
    def generate_from_fragment(fragment, sequence=None, dirn='.'):
        return ChunkReference(None)

    @staticmethod
    def read_reference_sequence_at_position_opened_file(f, start, end):
        """
        Read sequence from 1-indexed positions from reference.
        """
        f.seek(start - 1)
        sequence = f.read(end - start + 1).decode('utf-8')
        return sequence

    def read_reference_sequence_at_position(self, start, end):
        """
        Read sequence from 1-indexed positions from reference.
        """
        with gzip.open(self.ref_fn, 'rb') as f:
            f.seek(start - 1)
            sequence = f.read(end - start + 1).decode('utf-8')
        return sequence

    def __init__(self, ref_fn):
        self.ref_fn = ref_fn


class LocalChunkReference(ChunkReference):
    """
    Class for locally storing chunk sequences in reference files
    """

    @staticmethod
    def generate_from_fragment(fragment, sequence=None, dirn='.'):
        if sequence is None:
            sequence = fragment.indexed_fragment().sequence

        ref_fn = f"{dirn}/{fragment.id}.fa.gz"
        with gzip.open(ref_fn, 'wb') as f:
            sequence = f.write(sequence.encode('utf-8'))

        return LocalChunkReference(ref_fn)


class NetAppChunkReference(ChunkReference):
    """
    Class for storing chunk sequences in reference files on NetApp
    """

    @staticmethod
    def generate_from_fragment(fragment, sequence=None, dirn='.'):
        # NOTE: dirn is ignored in this method, as we're storing on NetApp
        if sequence is None:
            sequence = fragment.indexed_fragment().sequence
        ref_fn = fragment.build_fragment_fasta_from_sequence(sequence)

        return NetAppChunkReference(ref_fn)


class BigIntPrimaryModel(models.Model):
    class Meta:
        app_label = "edge"
        abstract = True

    id = models.BigIntegerField(primary_key=True)

    @transaction.atomic()
    def save(self, *args, **kwargs):
        # mimic auto_increment
        if self.id is None:
            klass = type(self)
            try:
                self.id = (
                    klass.objects.select_for_update()
                    .order_by("-id")
                    .values("id")[0]["id"]
                    + 1
                )
            except IndexError:
                self.id = 1
        return super(BigIntPrimaryModel, self).save(*args, **kwargs)


class Chunk(BigIntPrimaryModel):
    CHUNK_REFERENCE_CLASS = NetAppChunkReference

    class Meta:
        app_label = "edge"

    initial_fragment = models.ForeignKey("Fragment", on_delete=models.PROTECT)
    sequence = models.TextField(null=True)
    ref_start_index = models.PositiveIntegerField(blank=True, null=True)
    ref_end_index = models.PositiveIntegerField(blank=True, null=True)

    @property
    def ref_fn(self):
        return self.initial_fragment.fragment_reference_fasta_gz_fn()

    @property
    def is_sequence_based(self):
        return self.sequence is not None

    @property
    def is_reference_based(self):
        return None not in [self.ref_start_index, self.ref_end_index]

    @property
    def length(self):
        if self.is_sequence_based:
            return len(self.sequence)
        elif self.is_reference_based:
            return self.ref_end_index - self.ref_start_index + 1
        raise Exception("invalid chunk data")

    def get_sequence(self, f=None, start=None, end=None):
        if self.is_sequence_based:
            return self.sequence<commit_msg>fixed chunk ref_fn retrieval using id rather than pk<commit_after>import gzip
import json

from django.conf import settings
from django.db import models
from django.db import transaction


BULK_CREATE_BATCH_SIZE = 10000


class Annotation(object):
    """
    Can contain multiple Chunk_Feature objects merged together.
    """

    def __init__(self, base_first, base_last, chunk_feature, fragment=None):
        self.base_first = base_first
        self.base_last = base_last
        self.feature = chunk_feature.feature
        self.feature_base_first = chunk_feature.feature_base_first
        self.feature_base_last = chunk_feature.feature_base_last
        self.fragment = fragment

    @property
    def feature_name(self):
        if (
            self.feature_base_first != 1
            or self.feature_base_last != self.feature.length
        ):
            return "%s[%s:%s]" % (
                self.feature.name,
                self.feature_base_first,
                self.feature_base_last,
            )
        else:
            return self.feature.name

    def __str__(self):
        s = []
        s.append(self.feature_name)
        s.append(self.feature.type)
        if self.feature.strand == 1:
            s.append("+")
        else:
            s.append("-")
        return ", ".join(s)

    @staticmethod
    def from_chunk_feature_and_location_array(chunk_feature_locs):
        """
        chunk_feature_locs: an array of (Chunk_Feature, Fragment_Chunk_Location) tuples.
        """

        chunk_feature_locs = sorted(
            chunk_feature_locs,
            key=lambda t: (t[0].feature.id, t[1].base_first,
                           (t[0].feature.strand if t[0].feature.strand is not None else 1)
                           * t[0].feature_base_first)
        )

        annotations = []
        for cf, fcl in chunk_feature_locs:
            if (
                len(annotations) > 0
                and annotations[-1].feature.id == cf.feature_id
                and (((annotations[-1].feature.strand is None or annotations[-1].feature.strand > 0)
                      and annotations[-1].feature_base_last == cf.feature_base_first - 1)
                     or ((annotations[-1].feature.strand is not None
                          and annotations[-1].feature.strand < 0)
                         and annotations[-1].feature_base_first == cf.feature_base_last + 1))
                and annotations[-1].base_last == fcl.base_first - 1
            ):
                # merge annotation
                annotations[-1].base_last = fcl.base_last
                if (annotations[-1].feature.strand is None or annotations[-1].feature.strand > 0):
                    annotations[-1].feature_base_last = cf.feature_base_last
                else:
                    annotations[-1].feature_base_first = cf.feature_base_first
            else:
                annotations.append(
                    Annotation(
                        base_first=fcl.base_first,
                        base_last=fcl.base_last,
                        chunk_feature=cf,
                        fragment=fcl.fragment,
                    )
                )

        return sorted(annotations, key=lambda a: a.base_first)


class ChunkReference(object):
    """
    Parent class for storing chunk sequences in reference files
    1-based indexing example: (2, 4) of 5'-A(TCG)GCTA-3'
    """

    @staticmethod
    def generate_from_fragment(fragment, sequence=None, dirn='.'):
        return ChunkReference(None)

    @staticmethod
    def read_reference_sequence_at_position_opened_file(f, start, end):
        """
        Read sequence from 1-indexed positions from reference.
        """
        f.seek(start - 1)
        sequence = f.read(end - start + 1).decode('utf-8')
        return sequence

    def read_reference_sequence_at_position(self, start, end):
        """
        Read sequence from 1-indexed positions from reference.
        """
        with gzip.open(self.ref_fn, 'rb') as f:
            f.seek(start - 1)
            sequence = f.read(end - start + 1).decode('utf-8')
        return sequence

    def __init__(self, ref_fn):
        self.ref_fn = ref_fn


class LocalChunkReference(ChunkReference):
    """
    Class for locally storing chunk sequences in reference files
    """

    @staticmethod
    def generate_from_fragment(fragment, sequence=None, dirn='.'):
        if sequence is None:
            sequence = fragment.indexed_fragment().sequence

        ref_fn = f"{dirn}/{fragment.id}.fa.gz"
        with gzip.open(ref_fn, 'wb') as f:
            sequence = f.write(sequence.encode('utf-8'))

        return LocalChunkReference(ref_fn)


class NetAppChunkReference(ChunkReference):
    """
    Class for storing chunk sequences in reference files on NetApp
    """

    @staticmethod
    def generate_from_fragment(fragment, sequence=None, dirn='.'):
        # NOTE: dirn is ignored in this method, as we're storing on NetApp
        if sequence is None:
            sequence = fragment.indexed_fragment().sequence
        ref_fn = fragment.build_fragment_fasta_from_sequence(sequence)

        return NetAppChunkReference(ref_fn)


class BigIntPrimaryModel(models.Model):
    class Meta:
        app_label = "edge"
        abstract = True

    id = models.BigIntegerField(primary_key=True)

    @transaction.atomic()
    def save(self, *args, **kwargs):
        # mimic auto_increment
        if self.id is None:
            klass = type(self)
            try:
                self.id = (
                    klass.objects.select_for_update()
                    .order_by("-id")
                    .values("id")[0]["id"]
                    + 1
                )
            except IndexError:
                self.id = 1
        return super(BigIntPrimaryModel, self).save(*args, **kwargs)


class Chunk(BigIntPrimaryModel):
    CHUNK_REFERENCE_CLASS = NetAppChunkReference

    class Meta:
        app_label = "edge"

    initial_fragment = models.ForeignKey("Fragment", on_delete=models.PROTECT)
    sequence = models.TextField(null=True)
    ref_start_index = models.PositiveIntegerField(blank=True, null=True)
    ref_end_index = models.PositiveIntegerField(blank=True, null=True)

    @property
    def ref_fn(self):
        return f"{settings.SEQUENCE_FILE_DIR}/edge-fragment-{self.initial_fragment_id}.fa.gz"

    @property
    def is_sequence_based(self):
        return self.sequence is not None

    @property
    def is_reference_based(self):
        return None not in [self.ref_start_index, self.ref_end_index]

    @property
    def length(self):
        if self.is_sequence_based:
            return len(self.sequence)
        elif self.is_reference_based:
            return self.ref_end_index - self.ref_start_index + 1
        raise Exception("invalid chunk data")

    def get_sequence(self, f=None, start=None, end=None):
        if self.is_sequence_based:
            return self.sequence
{'fraction_non_alphanumeric': 0.09484565893154052, 'fraction_numerical': 0.007373671654738669, 'mean_word_length': 2.017888307155323, 'pattern_counts': {'":': 0, '<': 5, '<?xml version=': 0, '>': 9, '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': '3794918', 'n_tokens_mistral': 4037, 'n_tokens_neox': 3790, 'n_words': 839}