<commit_before>from django.contrib.postgres import lookups
from django.db import models


class PgLtreeField(models.TextField):

    description = 'Ltree'

    def db_type(self, connection):
        return 'ltree'

    def get_internal_type(self):
        return "CharField"

    def formfield(self, **kwargs):
        return super(models.TextField, self).formfield()


@PgLtreeField.register_lookup
class AncestorOf(lookups.PostgresSimpleLookup):
    lookup_name = 'ancestor_of'
    operator = '@>'


@PgLtreeField.register_lookup
class DescendantOf(lookups.PostgresSimpleLookup):
    lookup_name = 'descendant_of'
    operator = '<@'


@PgLtreeField.register_lookup
class Match(lookups.PostgresSimpleLookup):
    lookup_name = 'match'
    operator = '~'<commit_msg>add LtreeArray support<commit_after>from django.contrib.postgres import lookups
from django.contrib.postgres.fields.array import ArrayContainedBy, ArrayContains, \
    ArrayField
from django.db import models


class PgLtreeField(models.TextField):
    description = 'Ltree'
    
    def db_type(self, connection):
        return 'ltree'
    
    def get_internal_type(self):
        return "CharField"
    
    def get_prep_value(self, value):
        if isinstance(value, list) or isinstance(value, tuple):
            return [self.to_python(val) for val in value]
        return self.to_python(value)


class LtreeLookup(lookups.PostgresSimpleLookup):
    def as_sql(self, qn, connection):
        sql, params = super().as_sql(qn, connection)
        sql = '%s::%s' % (sql, self.lhs.output_field.db_type(connection))
        if isinstance(self.rhs, list) or isinstance(self.rhs, tuple):
            # Cast to array
            sql += '[]'
        return sql, params


@PgLtreeField.register_lookup
class AncestorOf(LtreeLookup):
    lookup_name = 'ancestor_of'
    operator = '@>'


@PgLtreeField.register_lookup
class DescendantOf(LtreeLookup):
    lookup_name = 'descendant_of'
    operator = '<@'


@PgLtreeField.register_lookup
class Match(LtreeLookup):
    lookup_name = 'match'
    operator = '~'


class LtreeArrayField(ArrayField):
    def __init__(self, **kwargs):
        kwargs['base_field'] = PgLtreeField()
        super().__init__(**kwargs)
    
    def get_db_prep_value(self, value, connection, prepared=False):
        if value:
            return '{%s}' % ','.join(value)
        return value
    
    def from_db_value(self, value, expression, connection, context):
        if value is None:
            return value
        # value example: '{1.2,1.3,1.4}'
        # remove brackets {}, then split by comma
        # return ['1.2', '1.3', '1.4']
        return value[1:-1].split(',')


@LtreeArrayField.register_lookup
class ArrayAncestorOf(ArrayContainedBy):
    lookup_name = 'ancestor_of'


@LtreeArrayField.register_lookup
class ArrayDescendantOf(ArrayContains):
    lookup_name = 'descendant_of'


@LtreeArrayField.register_lookup
class ArrayMatch(lookups.PostgresSimpleLookup):
    lookup_name = 'match'
    operator = '~'
    
    def as_sql(self, qn, connection):
        sql, params = super().as_sql(qn, connection)
        sql = '%s::%s' % (sql, self.lhs.output_field.db_type(connection))
        return sql, params
{'fraction_non_alphanumeric': 0.1135656502800249, 'fraction_numerical': 0.0043559427504667085, 'mean_word_length': 3.410150891632373, 'pattern_counts': {'":': 0, '<': 5, '<?xml version=': 0, '>': 5, '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': '900354', 'n_tokens_mistral': 1035, 'n_tokens_neox': 1014, 'n_words': 214}