root/branches/mbutscher/work/lib/whoosh/formats.py @ 234

#===============================================================================
# Copyright 2009 Matt Chaput
# 
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
# 
#    http://www.apache.org/licenses/LICENSE-2.0
# 
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#===============================================================================

"""
The classes in this module encode and decode posting information for a field.
The field format essentially determines what information is stored about each
occurance of a term.
"""

from collections import defaultdict
from struct import Struct
from cPickle import dumps, loads
from cStringIO import StringIO

from whoosh.analysis import unstopped
from whoosh.system import (_INT_SIZE, _FLOAT_SIZE, pack_uint, unpack_uint,
                           pack_float, unpack_float)
from whoosh.util import varint, read_varint, float_to_byte, byte_to_float


# Format base class

class Format(object):
    """Abstract base class representing a storage format for a field or vector.
    Format objects are responsible for writing and reading the low-level
    representation of a field. It controls what kind/level of information to
    store about the indexed fields.
    """
    
    posting_size = -1
    textual = True
    __inittypes__ = dict(analyzer=object, field_boost=float)
    
    def __init__(self, analyzer, field_boost=1.0, **options):
        """
        :param analyzer: The analysis.Analyzer object to use to index this
            field. See the analysis module for more information. If this value
            is None, the field is not indexed/searchable.
        :param field_boost: A constant boost factor to scale to the score
            of all queries matching terms in this field.
        """
        
        self.analyzer = analyzer
        self.field_boost = field_boost
        self.options = options
    
    def __eq__(self, other):
        return (other
                and self.__class__ is other.__class__
                and self.__dict__ == other.__dict__)
    
    def __repr__(self):
        return "%s(%r, boost = %s)" % (self.__class__.__name__,
                                       self.analyzer, self.field_boost)
    
    def clean(self):
        if self.analyzer and hasattr(self.analyzer, "clean"):
            self.analyzer.clean()
    
    def word_values(self, value, **kwargs):
        """Takes the text value to be indexed and yields a series of
        ("tokentext", frequency, weight, valuestring) tuples, where frequency
        is the number of times "tokentext" appeared in the value, weight is the
        weight (a float usually equal to frequency in the absence of per-term
        boosts) and valuestring is encoded field-specific posting value for the
        token. For example, in a Frequency format, the value string would be
        the same as frequency; in a Positions format, the value string would
        encode a list of token positions at which "tokentext" occured.
        
        :param value: The unicode text to index.
        """
        raise NotImplementedError
    
    def analyze(self, unicodestring, mode='', **kwargs):
        """Returns a :class:`whoosh.analysis.Token` iterator from the given
        unicode string.
        
        :param unicodestring: the string to analyzer.
        :param mode: a string indicating the purpose for which the unicode
            string is being analyzed, i.e. 'index' or 'query'.
        """
        
        if not self.analyzer:
            raise Exception("%s format has no analyzer" % self.__class__)
        return self.analyzer(unicodestring, mode=mode, **kwargs)
    
    def encode(self, value):
        """Returns the given value encoded as a string.
        """
        raise NotImplementedError
    
    def supports(self, name):
        """Returns True if this format supports interpreting its posting
        value as 'name' (e.g. "frequency" or "positions").
        """
        return hasattr(self, "decode_" + name)
    
    def decoder(self, name):
        """Returns the bound method for interpreting value as 'name',
        where 'name' is for example "frequency" or "positions". This
        object must have a corresponding Format.decode_<name>() method.
        """
        return getattr(self, "decode_" + name)
    
    def decode_as(self, astype, valuestring):
        """Interprets the encoded value string as 'astype', where 'astype' is
        for example "frequency" or "positions". This object must have a
        corresponding decode_<astype>() method.
        """
        return self.decoder(astype)(valuestring)
    

# Concrete field classes

class Existence(Format):
    """Only indexes whether a given term occurred in a given document; it does
    not store frequencies or positions. This is useful for fields that should
    be searchable but not scorable, such as file path.
    
    Supports: frequency, weight (always reports frequency = 1).
    """
    
    posting_size = 0
    __inittypes__ = dict(analyzer=object, field_boost=float)
    
    def __init__(self, analyzer, field_boost=1.0, **options):
        self.analyzer = analyzer
        self.field_boost = field_boost
        self.options = options
    
    def word_values(self, value, **kwargs):
        wordset = set(t.text for t
                      in unstopped(self.analyzer(value, **kwargs)))
        return ((w, 1, 1.0, '') for w in wordset)
    
    def encode(self, value):
        return ''
    
    def decode_frequency(self, valuestring):
        return 1
    
    def decode_weight(self, valuestring):
        return self.field_boost


class Frequency(Format):
    """Stores frequency information for each posting.
    
    Supports: frequency, weight.
    """
    
    posting_size = _INT_SIZE
    __inittypes__ = dict(analyzer=object, field_boost=float,
                         boost_as_freq=bool)
    
    def __init__(self, analyzer, field_boost=1.0, boost_as_freq=False,
                 **options):
        """
        :param analyzer: The analysis.Analyzer object to use to index this
            field. See the analysis module for more information. If this value
            is None, the field is not indexed/searchable.
        :param field_boost: A constant boost factor to scale to the score of
            all queries matching terms in this field.
        """
        
        self.analyzer = analyzer
        self.field_boost = field_boost
        self.options = options
        
    def word_values(self, value, **kwargs):
        freqs = defaultdict(int)
        weights = defaultdict(float)
        
        for t in unstopped(self.analyzer(value, boosts=True, **kwargs)):
            freqs[t.text] += 1
            weights[t.text] += t.boost
        
        encode = self.encode
        return ((w, freq, weights[w], encode(freq))
                for w, freq in freqs.iteritems())

    def encode(self, freq):
        return pack_uint(freq)
    
    def decode_frequency(self, valuestring):
        return unpack_uint(valuestring)[0]
    
    def decode_weight(self, valuestring):
        freq = unpack_uint(valuestring)[0]
        return freq * self.field_boost
    

class DocBoosts(Frequency):
    """A Field that stores frequency and per-document boost information for
    each posting.
    
    Supports: frequency, weight.
    """
    
    posting_size = _INT_SIZE + 1
    
    def word_values(self, value, doc_boost=1.0, **kwargs):
        freqs = defaultdict(int)
        weights = defaultdict(float)
        for t in unstopped(self.analyzer(value, boosts=True, **kwargs)):
            weights[t.text] += t.boost
            freqs[t.text] += 1
        
        encode = self.encode
        return ((w, freq, weights[w] * doc_boost, encode((freq, doc_boost)))
                for w, freq in freqs.iteritems())
    
    def encode(self, freq_docboost):
        freq, docboost = freq_docboost
        return pack_uint(freq) + float_to_byte(docboost)
    
    def decode_docboosts(self, valuestring):
        freq = unpack_uint(valuestring[:_INT_SIZE])[0]
        docboost = byte_to_float(valuestring[-1])
        return (freq, docboost)
    
    def decode_frequency(self, valuestring):
        return unpack_uint(valuestring[0:_INT_SIZE])[0]
    
    def decode_weight(self, valuestring):
        freq = unpack_uint(valuestring[:_INT_SIZE])[0]
        docboost = byte_to_float(valuestring[-1])
        return freq * docboost * self.field_boost
    

# Vector formats

class Positions(Format):
    """A vector that stores position information in each posting, to allow
    phrase searching and "near" queries.
    
    Supports: frequency, weight, positions, position_boosts (always reports
    position boost = 1.0).
    """
    
    def word_values(self, value, start_pos=0, **kwargs):
        poses = defaultdict(list)
        weights = defaultdict(float)
        for t in unstopped(self.analyzer(value, positions=True,
                                         start_pos=start_pos, **kwargs)):
            poses[t.text].append(start_pos + t.pos)
            weights[t.text] += t.boost
        
        encode = self.encode
        return ((w, len(poslist), weights[w], encode(poslist))
                for w, poslist in poses.iteritems())
    
    def encode(self, positions):
        codes = []
        base = 0
        for pos in positions:
            codes.append(pos - base)
            base = pos
        return pack_uint(len(codes)) + dumps(codes, -1)[2:-1]
    
    def decode_positions(self, valuestring):
        codes = loads(valuestring[_INT_SIZE:] + ".")
        position = 0
        positions = []
        for code in codes:
            position += code
            positions.append(position)
        return positions
    
    def decode_frequency(self, valuestring):
        return unpack_uint(valuestring[:_INT_SIZE])[0]
    
    def decode_weight(self, valuestring):
        return self.decode_frequency(valuestring) * self.field_boost
    
    def decode_position_boosts(self, valuestring):
        return [(pos, 1) for pos in self.decode_positions(valuestring)]
    

class Characters(Positions):
    """Stores token position and character start and end information for each
    posting.
    
    Supports: frequency, weight, positions, position_boosts (always reports
    position boost = 1.0), characters.
    """
    
    def word_values(self, value, start_pos=0, start_char=0, **kwargs):
        seen = defaultdict(list)
        
        for t in unstopped(self.analyzer(value, positions=True, chars=True,
                                         start_pos=start_pos,
                                         start_char=start_char, **kwargs)):
            seen[t.text].append((t.pos, start_char + t.startchar,
                                 start_char + t.endchar))
        
        encode = self.encode
        return ((w, len(ls), float(len(ls)), encode(ls))
                for w, ls in seen.iteritems())
    
    def encode(self, posns_chars):
        # posns_chars = [(pos, startchar, endchar), ...]
        codes = []
        posbase = 0
        charbase = 0
        for pos, startchar, endchar in posns_chars:
            codes.append((pos - posbase, startchar - charbase, endchar - startchar))
            posbase = pos
            charbase = endchar
        return pack_uint(len(posns_chars)) + dumps(codes, -1)[2:-1]
    
    def decode_characters(self, valuestring):
        codes = loads(valuestring[_INT_SIZE:] + ".")
        position = 0
        endchar = 0
        posns_chars = []
        for code in codes:
            position = code[0] + position
            startchar = code[1] + endchar
            endchar = code[2] + startchar
            posns_chars.append((position, startchar, endchar))
        return posns_chars
    
    def decode_positions(self, valuestring):
        codes = loads(valuestring[_INT_SIZE:] + ".")
        position = 0
        posns = []
        for code in codes:
            position = code[0] + position
            posns.append(position)
        return posns
    

class PositionBoosts(Positions):
    """A format that stores positions and per-position boost information
    in each posting.
    
    Supports: frequency, weight, positions, position_boosts.
    """
    
    def word_values(self, value, start_pos=0, **kwargs):
        seen = defaultdict(iter)
        for t in unstopped(self.analyzer(value, positions=True, boosts=True,
                                         start_pos=start_pos, **kwargs)):
            pos = t.pos
            boost = t.boost
            seen[t.text].append((pos, boost))
        
        encode = self.encode
        return ((w, len(poslist), sum(p[1] for p in poslist), encode(poslist))
                for w, poslist in seen.iteritems())
    
    def encode(self, posns_boosts):
        # posns_boosts = [(pos, boost), ...]
        codes = []
        base = 0
        summedboost = 0
        for pos, boost in posns_boosts:
            summedboost += boost
            codes.append((pos - base, boost))
            base = pos
            
        return (pack_uint(len(posns_boosts)) + pack_float(summedboost)
                + dumps(codes, -1)[2:-1])
        
    def decode_position_boosts(self, valuestring):
        codes = loads(valuestring[_INT_SIZE + _FLOAT_SIZE:] + ".")
        position = 0
        posns_boosts = []
        for code in codes:
            position = code[0] + position
            posns_boosts.append((position, code[1]))
        return posns_boosts
    
    def decode_positions(self, valuestring):
        codes = loads(valuestring[_INT_SIZE + _FLOAT_SIZE:] + ".")
        position = 0
        posns = []
        for code in codes:
            position = code[0] + position
            posns.append(position)
        return posns
    
    def decode_weight(self, valuestring):
        summedboost = unpack_float(valuestring[_INT_SIZE:_INT_SIZE + _FLOAT_SIZE])[0]
        return summedboost
    

class CharacterBoosts(Characters):
    """A format that stores positions, character start and end, and
    per-position boost information in each posting.
    
    Supports: frequency, weight, positions, position_boosts, characters,
    character_boosts.
    """
    
    def word_values(self, value, start_pos=0, start_char=0, **kwargs):
        seen = defaultdict(iter)
        for t in unstopped(self.analyzer(value, positions=True,
                                         characters=True, boosts=True,
                                         start_pos=start_pos,
                                         start_char=start_char, **kwargs)):
            seen[t.text].append((t.pos,
                                 start_char + t.startchar,
                                 start_char + t.endchar,
                                 t.boost))
        
        encode = self.encode
        return ((w, len(poslist), sum(p[3] for p in poslist), encode(poslist))
                for w, poslist in seen.iteritems())
    
    def encode(self, posns_chars_boosts):
        # posns_chars_boosts = [(pos, startchar, endchar, boost), ...]
        codes = []
        posbase = 0
        charbase = 0
        summedboost = 0
        for pos, startchar, endchar, boost in posns_chars_boosts:
            codes.append((pos - posbase, startchar - charbase,
                          endchar - startchar, boost))
            posbase = pos
            charbase = endchar
            summedboost += boost
        
        return (pack_uint(len(posns_chars_boosts)) + pack_float(summedboost)
                + dumps(codes, -1)[2:-1])
        
    def decode_character_boosts(self, valuestring):
        codes = loads(valuestring[_INT_SIZE + _FLOAT_SIZE:] + ".")
        position = 0
        endchar = 0
        posn_char_boosts = []
        for code in codes:
            position = position + code[0]
            startchar = endchar + code[1]
            endchar = startchar + code[2]
            posn_char_boosts.append((position, startchar, endchar, code[3]))
        return posn_char_boosts
    
    def decode_positions(self, valuestring):
        return [item[0] for item in self.decode_character_boosts(valuestring)]
    
    def decode_characters(self, valuestring):
        return [(pos, startchar, endchar) for pos, startchar, endchar, _
                in self.decode_character_boosts(valuestring)]
    
    def decode_position_boosts(self, valuestring):
        return [(pos, boost) for pos, _, _, boost
                in self.decode_character_boosts(valuestring)]