<commit_before>from django.db import models
from django.db.models.fields import FieldDoesNotExist
from django.db.models.sql.constants import JOIN_TYPE, LHS_ALIAS, LHS_JOIN_COL, \
    TABLE_NAME, RHS_JOIN_COL
from django.utils.tree import Node
from djangotoolbox.fields import ListField
from .lookups import StandardLookup

OR = 'OR'

# TODO: optimize code
class BaseResolver(object):
    def __init__(self):
        # mapping from lookups to indexes
        self.index_map = {}
        # mapping from column names to field names
        self.column_to_name = {}

    ''' API called by resolver'''

    def create_index(self, lookup):
        field_to_index = self.get_field_to_index(lookup.model, lookup.field_name)

        # backend doesn't now how to handle this index definition
        if not field_to_index:
            return

        index_field = lookup.get_field_to_add(field_to_index)
        config_field = index_field.item_field if \
            isinstance(index_field, ListField) else index_field
        if hasattr(field_to_index, 'max_length') and \
                isinstance(config_field, models.CharField):
            config_field.max_length = field_to_index.max_length

        # don't install a field if it already exists
        try:
            lookup.model._meta.get_field(self.index_name(lookup))
        except:
            lookup.model.add_to_class(self.index_name(lookup), index_field)
            self.index_map[lookup] = index_field
            self.add_column_to_name(lookup.model, lookup.field_name)
        else:
            # makes dbindexer unit test compatible
            if lookup not in self.index_map:
                self.index_map[lookup] = lookup.model._meta.get_field(
                    self.index_name(lookup))
                self.add_column_to_name(lookup.model, lookup.field_name)

    def convert_insert_query(self, query):
        '''Converts a database saving query.'''

        for lookup in self.index_map.keys():
            self._convert_insert_query(query, lookup)

    def _convert_insert_query(self, query, lookup):
        if not lookup.model == query.model:
            return

        values = self.get_value(lookup.model, lookup.field_name, query)
        for i, obj in enumerate(query.objs):
            value = lookup.convert_value(values[i])
            field = obj._meta.get_field(lookup.index_name)
            if field.rel:
                value = field.rel.to.objects.get(pk=value)
            setattr(obj, lookup.index_name, value)

    def convert_filters(self, query):
        self._convert_filters(query, query.where)

    ''' helper methods '''

    def _convert_filters(self, query, filters):
        for index, child in enumerate(filters.children[:]):
            if isinstance(child, Node):
                self._convert_filters(query, child)
                continue

            self.convert_filter(query, filters, child, index)

    def convert_filter(self, query, filters, child, index):
        constraint, lookup_type, annotation, value = child

        if constraint.field is None:
            return

        field_name = self.column_to_name.get(constraint.field.column)
        if field_name and constraint.alias == \
                query.table_map[query.model._meta.db_table][0]:
            for lookup in self.index_map.keys():
                if lookup.matches_filter(query.model, field_name, lookup_type,
                                         value):
                    new_lookup_type, new_value = lookup.convert_lookup(value,
                                                                       lookup_type)
                    index_name = self.index_name(lookup)
                    self._convert_filter(query, filters, child, index,
                                         new_lookup_type, new_value, index_name)

    def _convert_filter(self, query, filters, child, index, new_lookup_type,
                        new_value, index_name):
        constraint, lookup_type, annotation, value = child
        lookup_type, value = new_lookup_type, new_value
        constraint.field = query.get_meta().get_field(index_name)
        constraint.col = constraint.field.column
        child = constraint, lookup_type, annotation, value
        filters.children[index] = child

    def index_name(self, lookup):
        return lookup.index_name

    def get_field_to_index(self, model, field_name):
        try:
            return model._meta.get_field(field_name)
        except:
            return None

    def get_value(self, model, field_name, query):
        from django.db import router, connections

        field_to_index = self.get_field_to_index(model, field_name)

        if field_to_index not in query.fields:
            raise FieldDoesNotExist('Cannot find field in query.')

        values = []
        for obj in query.objs:
            value = field_to_index.get_db_prep_save(
                field_to_index.value_from_object(obj),
                connections[router.db_for_read(model)]
            )
            values.append(value)

        return values

    def add_column_to_name(self, model, field_name):
        column_name = model._meta.get_field(field_name).column
        self.column_to_name[column_name] = field_name

    def get_index(self, lookup):
        return self.index_map[lookup]

    def get_query_position(self, query, lookup):
        for index, field in enumerate(query.fields):
            if field is self.get_index(lookup):
                return index
        return None

def unref_alias(query, alias):
    table_name = query.alias_map[alias][TABLE_NAME]
    query.alias_refcount[alias] -= 1
    if query.alias_refcount[alias] < 1:
        # Remove all information about the join
        del query.alias_refcount[alias]
        del query.join_map[query.rev_join_map[alias]]
        del query.rev_join_map[alias]
        del query.alias_map[alias]
        query.tables.remove(alias)
        query.table_map[table_name].remove(alias)
        if len(query.table_map[table_name]) == 0:
            del query.table_map[table_name]
        query.used_aliases.discard(alias)

class FKNullFix(BaseResolver):
    '''
        Django doesn't generate correct code for ForeignKey__isnull.
        It becomes a JOIN with pk__isnull which won't work on nonrel DBs,
        so we rewrite the JOIN here.
    '''

    def create_index(self, lookup):
        pass

    def convert_insert_query(self, query):
        pass

    def convert_filter(self, query, filters, child, index):
        constraint, lookup_type, annotation, value = child
        if constraint.field is not None and lookup_type == 'isnull' and \
                        isinstance(constraint.field, models.ForeignKey):
            self.fix_fk_null_filter(query, constraint)

    def unref_alias(self, query, alias):
        unref_alias(query, alias)

    def fix_fk_null_filter(self, query, constraint):
        alias = constraint.alias<commit_msg>Use a better test to find out if a field is a related field<commit_after>from django.db import models
from django.db.models.fields import FieldDoesNotExist
from django.db.models.sql.constants import JOIN_TYPE, LHS_ALIAS, LHS_JOIN_COL, \
    TABLE_NAME, RHS_JOIN_COL
from django.utils.tree import Node
from djangotoolbox.fields import ListField
from .lookups import StandardLookup

OR = 'OR'

# TODO: optimize code
class BaseResolver(object):
    def __init__(self):
        # mapping from lookups to indexes
        self.index_map = {}
        # mapping from column names to field names
        self.column_to_name = {}

    ''' API called by resolver'''

    def create_index(self, lookup):
        field_to_index = self.get_field_to_index(lookup.model, lookup.field_name)

        # backend doesn't now how to handle this index definition
        if not field_to_index:
            return

        index_field = lookup.get_field_to_add(field_to_index)
        config_field = index_field.item_field if \
            isinstance(index_field, ListField) else index_field
        if hasattr(field_to_index, 'max_length') and \
                isinstance(config_field, models.CharField):
            config_field.max_length = field_to_index.max_length

        # don't install a field if it already exists
        try:
            lookup.model._meta.get_field(self.index_name(lookup))
        except:
            lookup.model.add_to_class(self.index_name(lookup), index_field)
            self.index_map[lookup] = index_field
            self.add_column_to_name(lookup.model, lookup.field_name)
        else:
            # makes dbindexer unit test compatible
            if lookup not in self.index_map:
                self.index_map[lookup] = lookup.model._meta.get_field(
                    self.index_name(lookup))
                self.add_column_to_name(lookup.model, lookup.field_name)

    def convert_insert_query(self, query):
        '''Converts a database saving query.'''

        for lookup in self.index_map.keys():
            self._convert_insert_query(query, lookup)

    def _convert_insert_query(self, query, lookup):
        if not lookup.model == query.model:
            return

        values = self.get_value(lookup.model, lookup.field_name, query)
        for i, obj in enumerate(query.objs):
            value = lookup.convert_value(values[i])
            field = obj._meta.get_field(lookup.index_name)
            if field.rel:
                value = field.rel.to.objects.get(pk=value)
            setattr(obj, lookup.index_name, value)

    def convert_filters(self, query):
        self._convert_filters(query, query.where)

    ''' helper methods '''

    def _convert_filters(self, query, filters):
        for index, child in enumerate(filters.children[:]):
            if isinstance(child, Node):
                self._convert_filters(query, child)
                continue

            self.convert_filter(query, filters, child, index)

    def convert_filter(self, query, filters, child, index):
        constraint, lookup_type, annotation, value = child

        if constraint.field is None:
            return

        field_name = self.column_to_name.get(constraint.field.column)
        if field_name and constraint.alias == \
                query.table_map[query.model._meta.db_table][0]:
            for lookup in self.index_map.keys():
                if lookup.matches_filter(query.model, field_name, lookup_type,
                                         value):
                    new_lookup_type, new_value = lookup.convert_lookup(value,
                                                                       lookup_type)
                    index_name = self.index_name(lookup)
                    self._convert_filter(query, filters, child, index,
                                         new_lookup_type, new_value, index_name)

    def _convert_filter(self, query, filters, child, index, new_lookup_type,
                        new_value, index_name):
        constraint, lookup_type, annotation, value = child
        lookup_type, value = new_lookup_type, new_value
        constraint.field = query.get_meta().get_field(index_name)
        constraint.col = constraint.field.column
        child = constraint, lookup_type, annotation, value
        filters.children[index] = child

    def index_name(self, lookup):
        return lookup.index_name

    def get_field_to_index(self, model, field_name):
        try:
            return model._meta.get_field(field_name)
        except:
            return None

    def get_value(self, model, field_name, query):
        from django.db import router, connections

        field_to_index = self.get_field_to_index(model, field_name)

        if field_to_index not in query.fields:
            raise FieldDoesNotExist('Cannot find field in query.')

        values = []
        for obj in query.objs:
            value = field_to_index.get_db_prep_save(
                field_to_index.value_from_object(obj),
                connections[router.db_for_read(model)]
            )
            values.append(value)

        return values

    def add_column_to_name(self, model, field_name):
        column_name = model._meta.get_field(field_name).column
        self.column_to_name[column_name] = field_name

    def get_index(self, lookup):
        return self.index_map[lookup]

    def get_query_position(self, query, lookup):
        for index, field in enumerate(query.fields):
            if field is self.get_index(lookup):
                return index
        return None

def unref_alias(query, alias):
    table_name = query.alias_map[alias][TABLE_NAME]
    query.alias_refcount[alias] -= 1
    if query.alias_refcount[alias] < 1:
        # Remove all information about the join
        del query.alias_refcount[alias]
        del query.join_map[query.rev_join_map[alias]]
        del query.rev_join_map[alias]
        del query.alias_map[alias]
        query.tables.remove(alias)
        query.table_map[table_name].remove(alias)
        if len(query.table_map[table_name]) == 0:
            del query.table_map[table_name]
        query.used_aliases.discard(alias)

class FKNullFix(BaseResolver):
    '''
        Django doesn't generate correct code for ForeignKey__isnull.
        It becomes a JOIN with pk__isnull which won't work on nonrel DBs,
        so we rewrite the JOIN here.
    '''

    def create_index(self, lookup):
        pass

    def convert_insert_query(self, query):
        pass

    def convert_filter(self, query, filters, child, index):
        constraint, lookup_type, annotation, value = child
        if constraint.field is not None and lookup_type == 'isnull' and \
                        getattr(constraint.field, 'rel', None):
            self.fix_fk_null_filter(query, constraint)

    def unref_alias(self, query, alias):
        unref_alias(query, alias)

    def fix_fk_null_filter(self, query, constraint):
        alias = constraint.alias
{'fraction_non_alphanumeric': 0.08186714542190306, 'fraction_numerical': 0.0005745062836624775, 'mean_word_length': 2.3236276849642006, 'pattern_counts': {'":': 0, '<': 5, '<?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': 14, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '2602720', 'n_tokens_mistral': 3969, 'n_tokens_neox': 3831, 'n_words': 915}