<commit_before>from __future__ import unicode_literals

import json
import decimal

from django.core.exceptions import ValidationError
from django.db import models
from django.utils import six

from psycopg2.extras import register_json

# Actually register jsonb!
# http://schinckel.net/2014/05/24/python%2C-postgres-and-jsonb/
register_json(oid=3802, array_oid=3807)


    

class JSONField(six.with_metaclass(models.SubfieldBase, models.Field)):
    description = 'JSON Field'
    
    def __init__(self, *args, **kwargs):
        self.decode_kwargs = kwargs.pop('decode_kwargs', {
            'parse_float': decimal.Decimal
        })
        self.encode_kwargs = kwargs.pop('encode_kwargs', {
            'default': default
        })
        super(JSONField, self).__init__(*args, **kwargs)
    
    def get_internal_type(self):
        return 'JSONField'
    
    def db_type(self, connection):
        return 'jsonb'

    def get_db_prep_value(self, value, connection=None, prepared=None):
        return self.get_prep_value(value)
    
    def get_prep_value(self, value):
        if value is None:
            if not self.null and self.blank:
                return ""
            return None
        return json.dumps(value, **self.encode_kwargs)
    
    def get_prep_lookup(self, lookup_type, value, prepared=False):
        if lookup_type == 'has_key':
            # Need to ensure we have a string, as no other
            # values is appropriate.
            if not isinstance(value, six.string_types):
                value = '%s' % value
        if lookup_type == 'has_keys':
            if isinstance(value, six.string_types):
                value = [value]
            # This will cast numbers to strings, but also grab the keys
            # from a dict.
            value = ['%s' % v for v in value]
                
        return value
    
    def get_db_prep_lookup(self, lookup_type, value, connection, prepared=False):
        if lookup_type == 'contains':
            value = self.get_prep_value(value)
            return [value]
        
        return super(JSONField, self).get_db_prep_lookup(lookup_type, value, connection, prepared)
        
    def to_python(self, value):
        if isinstance(value, six.string_types) or value is None:
            if not value:
                if self.blank:
                    return ''
                return None
            try:
                return json.loads(value, **self.decode_kwargs)
            except ValueError as exc:
                raise ValidationError(*exc.args)
        
        # Should we round-trip a value when it is assigned?
        # It means we will always have data in the format that
        # it will come back from the server, which is a good
        # thing.
        return self.to_python(self.get_prep_value(value))


from django.db.models import lookups, Lookup


class HasKey(Lookup):
    lookup_name = 'has_key'
    
    def as_sql(self, qn, connection):
        lhs, lhs_params = self.process_lhs(qn, connection)
        rhs, rhs_params = self.process_rhs(qn, connection)
        params = lhs_params + rhs_params
        
        return '%s ? %s' % (lhs, rhs), params

JSONField.register_lookup(HasKey)


class Contains(Lookup):
    lookup_name = 'contains'
    
    def as_sql(self, qn, connection):
        lhs, lhs_params = self.process_lhs(qn, connection)
        rhs, rhs_params = self.process_rhs(qn, connection)
        params = lhs_params + rhs_params
        
        return '%s @> %s' % (lhs, rhs), params

JSONField.register_lookup(Contains)


# Other lookups:

# in            <@
# all_keys      ?&
# any_keys      ?|

def default(o):
    if hasattr(o, 'to_json'):
        return o.to_json()
    if isinstance(o, Decimal):
        return str(o)<commit_msg>Implement other lookup types for jsonb.<commit_after>from __future__ import unicode_literals

import json
import decimal

from django.core.exceptions import ValidationError
from django.db import models
from django.utils import six

from psycopg2.extras import register_json

# Actually register jsonb!
# http://schinckel.net/2014/05/24/python%2C-postgres-and-jsonb/
register_json(oid=3802, array_oid=3807)


    

class JSONField(six.with_metaclass(models.SubfieldBase, models.Field)):
    description = 'JSON Field'
    
    def __init__(self, *args, **kwargs):
        self.decode_kwargs = kwargs.pop('decode_kwargs', {
            'parse_float': decimal.Decimal
        })
        self.encode_kwargs = kwargs.pop('encode_kwargs', {
            'default': default
        })
        super(JSONField, self).__init__(*args, **kwargs)
    
    def get_internal_type(self):
        return 'JSONField'
    
    def db_type(self, connection):
        return 'jsonb'

    def get_db_prep_value(self, value, connection=None, prepared=None):
        return self.get_prep_value(value)
    
    def get_prep_value(self, value):
        if value is None:
            if not self.null and self.blank:
                return ""
            return None
        return json.dumps(value, **self.encode_kwargs)
    
    def get_prep_lookup(self, lookup_type, value, prepared=False):
        if lookup_type == 'has_key':
            # Need to ensure we have a string, as no other
            # values is appropriate.
            if not isinstance(value, six.string_types):
                value = '%s' % value
        if lookup_type == 'has_keys':
            if isinstance(value, six.string_types):
                value = [value]
            # This will cast numbers to strings, but also grab the keys
            # from a dict.
            value = ['%s' % v for v in value]
                
        return value
    
    def get_db_prep_lookup(self, lookup_type, value, connection, prepared=False):
        if lookup_type == 'contains':
            value = self.get_prep_value(value)
            return [value]
        
        return super(JSONField, self).get_db_prep_lookup(lookup_type, value, connection, prepared)
        
    def to_python(self, value):
        if isinstance(value, six.string_types) or value is None:
            if not value:
                if self.blank:
                    return ''
                return None
            try:
                return json.loads(value, **self.decode_kwargs)
            except ValueError as exc:
                raise ValidationError(*exc.args)
        
        # Should we round-trip a value when it is assigned?
        # It means we will always have data in the format that
        # it will come back from the server, which is a good
        # thing.
        return self.to_python(self.get_prep_value(value))


from django.db.models.lookups import BuiltinLookup, Lookup

class PostgresLookup(BuiltinLookup):
    def process_lhs(self, qn, connection, lhs=None):
        lhs = lhs or self.lhs
        return qn.compile(lhs)
    
    def get_rhs_op(self, connection, rhs):
        return '%s %s' % (self.operator, rhs)

class HasKey(PostgresLookup):
    lookup_name = 'has_key'
    operator = '?'

JSONField.register_lookup(HasKey)


class Contains(PostgresLookup):
    lookup_name = 'contains'
    operator = '@>'
    
JSONField.register_lookup(Contains)

class In(PostgresLookup):
    lookup_name = 'in'
    operator = '<@'

JSONField.register_lookup(In)

class AllKeys(PostgresLookup):
    lookup_name = 'all_keys'
    operator = '?&'

JSONField.register_lookup(AllKeys)

class AnyKeys(PostgresLookup):
    lookup_name = 'any_keys'
    operator = '?|'

JSONField.register_lookup(AnyKeys)


def default(o):
    if hasattr(o, 'to_json'):
        return o.to_json()
    if isinstance(o, Decimal):
        return str(o)
{'fraction_non_alphanumeric': 0.09390530996599529, 'fraction_numerical': 0.004708344232278315, 'mean_word_length': 2.297542043984476, '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': '5154557', 'n_tokens_mistral': 2305, 'n_tokens_neox': 2244, 'n_words': 597}