<commit_before>import time
from curious import model_registry
from curious.graph import traverse, mk_filter_function
from .parser import Parser
from .utils import report_time


class Query(object):

  def __init__(self, query):
    parser = Parser(query)
    self.__query = query
    self.__obj_query = parser.object_query
    self.__steps = parser.steps
    self.__validate()


  @property
  def query_string(self):
    return self.__query


  @staticmethod
  def _validate(query):
    """
    Validate a query. A query is an array whose elements are model
    relationships or subqueries. This function checks each model relationship
    to make sure the model and the relationship exist.
    """
    
    for rel in query:
      if 'orquery' in rel:
        for q in rel['orquery']:
          Query._validate(q)
      elif 'subquery' in rel:
        Query._validate(rel['subquery'])
      else:
        model = rel['model']
        method = rel['method']
        if method is None:
          model_registry.get_manager(model).model_class
        else:
          model_registry.get_manager(model).getattr(method)


  def __validate(self):
    Query._validate([self.__obj_query]+self.__steps)


  def __get_objects(self):
    """
    Get initial objects from object query.
    """

    model = self.__obj_query['model']
    method = self.__obj_query['method']
    filters = self.__obj_query['filters']
    filter_f = mk_filter_function(filters)

    if method is None:
      cls = model_registry.get_manager(model).model_class
      q = cls.objects.all()
      q = filter_f(q)
      return q
    else:
      f = model_registry.get_manager(model).getattr(method)
      return f(filter_f)


  @staticmethod
  def _extend_result(obj_src, next_obj_src):
    # build input hash of IDs
    input_map = {}
    for obj, src in obj_src:
      if obj.pk not in input_map:
        input_map[obj.pk] = []
      input_map[obj.pk].append(src)

    keep = []
    for next_obj, next_src in next_obj_src:
      if next_src in input_map:
        for src in input_map[next_src]:
          keep.append((next_obj, src))

    return list(set(keep))


  @staticmethod
  @report_time
  def _graph_step(obj_src, model, step_f, filters):
    """
    Traverse one step on the graph. Takes in and returns arrays of output,
    input object tuples. The input objects in the tuples are from start of the
    query, not start of this step.
    """

    # check if type matches existing object type
    if len(obj_src):
      t = type(obj_src[0][0])
      if t._deferred:
        t = t.__base__
      if t != model_registry.get_manager(model).model_class:
        raise Exception('Type mismatch when executing query: expecting "%s", got "%s"' %
                        (model, type(obj_src[0][0])))

    next_obj_src = traverse([obj for obj, src in obj_src], step_f, filters)

    return Query._extend_result(obj_src, next_obj_src)


  @staticmethod
  def _recursive_rel(obj_src, step):
    """
    Traverse a relationship recursively. Collected objects, either loop
    terminating objects or loop continuing objects. Returns arrays of output,
    input object tuples. The input objects in the tuples are from start of the
    query, not start of this step.
    """

    model = step['model']
    method = step['method']
    filters = step['filters']
    collect = step['collect']
    step_f = model_registry.get_manager(model).getattr(method)

    collected = {}<commit_msg>Allow relationships from custom models<commit_after>import time
from curious import model_registry
from curious.graph import traverse, mk_filter_function
from .parser import Parser
from .utils import report_time


class Query(object):

  def __init__(self, query):
    parser = Parser(query)
    self.__query = query
    self.__obj_query = parser.object_query
    self.__steps = parser.steps
    self.__validate()


  @property
  def query_string(self):
    return self.__query


  @staticmethod
  def _validate(query):
    """
    Validate a query. A query is an array whose elements are model
    relationships or subqueries. This function checks each model relationship
    to make sure the model and the relationship exist.
    """
    
    for rel in query:
      if 'orquery' in rel:
        for q in rel['orquery']:
          Query._validate(q)
      elif 'subquery' in rel:
        Query._validate(rel['subquery'])
      else:
        model = rel['model']
        method = rel['method']
        if method is None:
          model_registry.get_manager(model).model_class
        else:
          model_registry.get_manager(model).getattr(method)


  def __validate(self):
    Query._validate([self.__obj_query]+self.__steps)


  def __get_objects(self):
    """
    Get initial objects from object query.
    """

    model = self.__obj_query['model']
    method = self.__obj_query['method']
    filters = self.__obj_query['filters']
    filter_f = mk_filter_function(filters)

    if method is None:
      cls = model_registry.get_manager(model).model_class
      q = cls.objects.all()
      q = filter_f(q)
      return q
    else:
      f = model_registry.get_manager(model).getattr(method)
      return f(filter_f)


  @staticmethod
  def _extend_result(obj_src, next_obj_src):
    # build input hash of IDs
    input_map = {}
    for obj, src in obj_src:
      if obj.pk not in input_map:
        input_map[obj.pk] = []
      input_map[obj.pk].append(src)

    keep = []
    for next_obj, next_src in next_obj_src:
      if next_src in input_map:
        for src in input_map[next_src]:
          keep.append((next_obj, src))

    return list(set(keep))


  @staticmethod
  @report_time
  def _graph_step(obj_src, model, step_f, filters):
    """
    Traverse one step on the graph. Takes in and returns arrays of output,
    input object tuples. The input objects in the tuples are from start of the
    query, not start of this step.
    """

    # check if type matches existing object type
    if len(obj_src):
      t = type(obj_src[0][0])
      if hasattr(t, '_deferred') and t._deferred:
        t = t.__base__
      if t != model_registry.get_manager(model).model_class:
        raise Exception('Type mismatch when executing query: expecting "%s", got "%s"' %
                        (model, type(obj_src[0][0])))

    next_obj_src = traverse([obj for obj, src in obj_src], step_f, filters)

    return Query._extend_result(obj_src, next_obj_src)


  @staticmethod
  def _recursive_rel(obj_src, step):
    """
    Traverse a relationship recursively. Collected objects, either loop
    terminating objects or loop continuing objects. Returns arrays of output,
    input object tuples. The input objects in the tuples are from start of the
    query, not start of this step.
    """

    model = step['model']
    method = step['method']
    filters = step['filters']
    collect = step['collect']
    step_f = model_registry.get_manager(model).getattr(method)

    collected = {}
{'fraction_non_alphanumeric': 0.0940515627250468, 'fraction_numerical': 0.001152239665850497, 'mean_word_length': 2.870680044593088, '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': '1253404', 'n_tokens_mistral': 2115, 'n_tokens_neox': 2056, 'n_words': 631}