1 # -*- coding: utf-8 -*-

     2 '''

     3 Created on Aug 1, 2012

     4 

     5 @author: ymh

     6 '''

     7 

     8 from django.conf import settings

     9 from haystack.query import SQ

    10 from whoosh.qparser import (SimpleParser, FieldsPlugin, OperatorsPlugin, 

    11     PhrasePlugin, SingleQuotePlugin, GroupPlugin, PrefixPlugin, GtLtPlugin, 

    12     RangePlugin)

    13 from whoosh.query import (Term, And, AndMaybe, Or, AndNot, Not, Phrase, Prefix, 

    14     TermRange)

    15 

    16 HAYSTACK_DEFAULT_OPERATOR = getattr(settings,'HAYSTACK_DEFAULT_OPERATOR','AND')

    17 

    18 class QueryParser(object):

    19 

    20 

    21     def __init__(self, fieldname):

    22         '''

    23         Constructor

    24         '''

    25         self.w_parser = SimpleParser(fieldname, None)

    26         self.w_parser.add_plugin(FieldsPlugin())

    27         self.w_parser.add_plugin(OperatorsPlugin())

    28         self.w_parser.add_plugin(PhrasePlugin())

    29         self.w_parser.add_plugin(SingleQuotePlugin())

    30         self.w_parser.add_plugin(GroupPlugin())

    31         self.w_parser.add_plugin(PrefixPlugin())

    32         self.w_parser.add_plugin(GtLtPlugin())

    33         self.w_parser.add_plugin(RangePlugin())

    34         self.query = None

    35         self.current_node_stack = []        

    36         

    37     def parse(self, query):

    38         

    39         self.query = SQ()

    40         self.current_node_stack = [(self.query, HAYSTACK_DEFAULT_OPERATOR)]

    41 

    42         wquery = self.w_parser.parse(query)

    43         

    44         self.visit(wquery)

    45         

    46         if len(self.query) == 1 and isinstance(self.query.children[0], SQ):

    47             return self.query.children[0]

    48         else:

    49             return self.query 

    50         

    51         

    52     def visit(self, q):

    53         

    54         if isinstance(q, Term):

    55             current_node, current_connector = self.current_node_stack.pop() 

    56             current_node.add(SQ(**{q.fieldname:q.text}), current_connector)

    57             self.current_node_stack.append((current_node,current_connector))

    58         elif isinstance(q, And):

    59             self._add_compound_query(q, SQ.AND)

    60         elif isinstance(q, AndMaybe):

    61             self._add_andmaybe(q)

    62         elif isinstance(q, Or):

    63             self._add_compound_query(q, SQ.OR)

    64         elif isinstance(q, AndNot):

    65             self._add_andnot(q)

    66         elif isinstance(q, Not):

    67             self._add_not(q)

    68         elif isinstance(q, Phrase):

    69             self._add_phrase(q)

    70         elif isinstance(q, Prefix):

    71             self._add_prefix(q)

    72         elif isinstance(q, TermRange):

    73             self._add_range(q)

    74             

    75     def _add_compound_query(self, q, connector):

    76 

    77         new_node = SQ()

    78         self.current_node_stack.append((new_node, connector))

    79         for subquery in q.subqueries:

    80             self.visit(subquery)

    81         self.current_node_stack.pop()

    82                         

    83         if len(new_node)==1 and isinstance(new_node.children[0], SQ) :

    84             new_node = new_node.children[0]

    85         

    86         current_node, current_connector = self.current_node_stack[-1]

    87         current_node.add(new_node, current_connector)

    88         

    89         

    90     def _add_andnot(self, q):

    91         

    92         new_node = SQ()

    93         self.current_node_stack.append((new_node, SQ.AND))

    94         self.visit(q.a)

    95         self.visit(Not(q.b))

    96         self.current_node_stack.pop()

    97         

    98         if len(new_node)==1 and isinstance(new_node.children[0], SQ) :

    99             new_node = new_node.children[0]

   100         

   101         current_node, current_connector = self.current_node_stack[-1]

   102         current_node.add(new_node, current_connector)

   103 

   104     def _add_andmaybe(self, q):

   105         

   106         new_node = SQ()

   107         self.current_node_stack.append((new_node, SQ.AND))

   108         self.visit(q.a)

   109         self.visit(q.b)

   110         self.current_node_stack.pop()

   111         

   112         if len(new_node)==1 and isinstance(new_node.children[0], SQ) :

   113             new_node = new_node.children[0]

   114         

   115         current_node, current_connector = self.current_node_stack[-1]

   116         current_node.add(new_node, current_connector)

   117 

   118         

   119     def _add_not(self, q):

   120         

   121         new_node = SQ()

   122         self.current_node_stack.append((new_node, SQ.AND))

   123         self.visit(q.query)

   124         self.current_node_stack.pop()

   125         

   126         if len(new_node)==1 and isinstance(new_node.children[0], SQ) :

   127             new_node = new_node.children[0]

   128             

   129         current_node, current_connector = self.current_node_stack[-1]

   130         current_node.add(~new_node, current_connector)

   131         

   132     def _add_phrase(self, q):

   133         new_node = SQ(**{q.fieldname+"__exact":" ".join(q.words)})            

   134         current_node, current_connector = self.current_node_stack[-1]

   135         current_node.add(new_node, current_connector)

   136 

   137     def _add_prefix(self, q):

   138         new_node = SQ(**{q.fieldname+"__startswith":q.text})            

   139         current_node, current_connector = self.current_node_stack[-1]

   140         current_node.add(new_node, current_connector)

   141 

   142     def _add_range(self, q):

   143         

   144         if q.start is None:

   145             if q.endexcl:

   146                 postfix = "__lt"

   147             else:

   148                 postfix = "__lte"

   149             new_node = SQ(**{q.fieldname+postfix:self.__convert_nb(q.end)})

   150         elif q.end is None:

   151             if q.startexcl:

   152                 postfix = "__gt"

   153             else:

   154                 postfix = "__gte"

   155             new_node = SQ(**{q.fieldname+postfix:self.__convert_nb(q.start)})

   156         else:

   157             new_node = SQ(**{q.fieldname+"__range":[self.__convert_nb(q.start),self.__convert_nb(q.end)]})

   158         

   159         current_node, current_connector = self.current_node_stack[-1]

   160         current_node.add(new_node, current_connector)

   161 

   162     def __convert_nb(self, str_nb):        

   163         try:

   164             res = int(str_nb)

   165             return res

   166         except ValueError:

   167             try:

   168                 res = float(str_nb)

   169                 return res

   170             except ValueError:

   171                 return str_nb

   172         

   173         

   174