root/tags/amplee-0.5.0/amplee/indexer/__init__.py

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

__doc__ = """Really simple indexing implementation for amplee

I assume something bigger such as Lucene would be much more efficient.

For example:

>>> import time
>>> from amplee.indexer import *
>>> ind = Indexer()
>>> container = ShelveContainer('/tmp/cache.p')
>>> pi = PublishedIndex('pi', container=container, granularity=DateIndex.month)
>>> container2 = ShelveContainer('/tmp/cache2.p')
>>> ai = AuthorIndex('ai', container=container2)
>>> ind.register(pi)
>>> ind.register(ai)

# With the indexer daemon
>>> ind.start_daemon(interval=1.0)
>>> try:
>>>     index(member)
>>>     time.sleep(3)
>>> finally:
>>>     ind.stop_daemon()

# Or without
>>> ind.process(member)

>>> from datetime import datetime
>>> start = datetime(2006, 10, 8, 14, 00, 12)
>>> end = datetime(2008, 10, 8, 16, 00, 12)
>>> r0 = pi.between(start, end)
>>> print r0 
>>> r1 = ai.lookup('Jon Doe')
>>> print r1

# You can even concatenate both set of results into one
>>> print r0 & r1

# Or substract (assuming r2)
>>> print r0 & r2 - r1

Once you get a result dictionnary you can do something like this:

>>> collection_name, member_id = r.pop()
>>> c0 = service.get_collection(collection_name)
>>> member = c0.get_member(member_id)

Or to get a list of member instances.

>>> items = ind.to_dict(r)
>>> for collection_name in items:
        c = service.get_collection(collection_name)
        members = c.reload_members_from_list(items[collection_name])
"""

from Queue import Queue, Empty
import sha
import threading
from datetime import datetime
import re
import shelve
from sets import Set

from amplee.utils import parse_isodate

__all__ = ['Indexer', 'BaseIndex', 'PublishedIndex', 'AuthorIndex',
           'CategoryIndex', 'UpdatedIndex', 'DateIndex', 'MemoryContainer',
           'index', 'ShelveContainer', 'MemcacheContainer',
           'KeywordIndex', 'AtomIDIndexer']

_queue = Queue()

def index(member):
    """Add the member to queue of elements to be indexed.

    The ``member`` is an instance of any of the member classes
    provided by amplee.

    Note that qe-queuing will be asynchronized and thus this
    function will return before the member is actually indexed.
    """
    _queue.put(member)

class MemoryContainer(object):
    """Dummy memory container based on the dictionnary.
    This does not do thread locking so be careful
    of what you do with it :)
    """
    def __init__(self):
        self._container = {}

    def shutdown(self):
        pass

    def __contains__(self, key):
        return key in self._container
        
    def __getitem__(self, key):
        if key in self._container:
            return self._container[key]
        raise KeyError
    
    def __setitem__(self, key, value):
        self._container[key] = value

    def __delitem__(self, key):
        if key in self._container:
            del self._container[key]

    def __iter__(self):
        return iter(self._container)
    
    def iterkeys(self):
        return self._container.iterkeys()
    
class ShelveContainer(object):
    def __init__(self, abs_path, protocol=2):
        """Simple container around the built-in shelve module.
        This provides a simple and more-or-less efficient
        persistence system.

        The ``abs_path`` is the absolute path of the shelve to
        open (and create if it does not exist).

        The ``protocol`` is the protocol value used by the open()
        function of the shelve module.

        If you plan to create your own container, they must
        implement the start and shutdown methods as-is.
        """
        self.abs_path = abs_path
        self.protocol = protocol
        self.shelf = shelve.open(self.abs_path, protocol=self.protocol,
                                 writeback=True)

    def shutdown(self):
        """Closes the shelve container when not needed anymore."""
        self.shelf.close()

    def __contains__(self, key):
        return key in self.shelf
        
    def __getitem__(self, key):
        if key in self.shelf:
            return self.shelf[key]
        raise KeyError
    
    def __setitem__(self, key, value):
        self.shelf[key] = value
        self.shelf.sync()

    def __delitem__(self, key):
        if key in self.shelf:
            del self.shelf[key]
        self.shelf.sync()

    def __iter__(self):
        return iter(self.shelf)
    
    def iterkeys(self):
        return self.shelf.iterkeys()
    
class MemcacheContainer(object):
    def __init__(self, servers, cache_key=None):
        """
        Container based on the memcached tool for highly
        efficient and distributed memory cache.

        To use this container you must have either
        cmemcache or python-memcache installed.

        The ``servers`` arguments is a list of bye string
        (not unicode) respecting the format of memcached
        connection strings.

        The ``cache_key`` is a unique value that will be used
        to track the indexes within memcache. If not
        provided one random string will created.

        Basically only one key will be created within memcache
        and a list will be associated as a value.

        Indexes will be stored in that list.
        """
        try:
            from cmemcache import Client
        except ImportError:
            from memcache import Client

        if not cache_key:
            import string, random
            CHARS = string.letters + string.octdigits + '_-'
            cache_key = ''.join(random.sample(CHARS, 25))
            
        self.cache_key = cache_key
        self.mc = Client(servers)
        self.mc.set(self.cache_key, {})

    def shutdown(self):
        """Closes down all connections to memcached nodes."""
        self.mc.disconnect_all()

    def __contains__(self, key):
        l = self.mc.get(self.cache_key)
        return key in l
        
    def __getitem__(self, key):
        l = self.mc.get(self.cache_key)
        if key in l:
            return l[key]
        raise KeyError
    
    def __setitem__(self, key, value):
        l = self.mc.get(self.cache_key)
        l[key] = value
        self.mc.set(key, l)

    def __delitem__(self, key):
        l = self.mc.get(self.cache_key)
        if key in l:
            del l[key]
        self.mc.set(key, l)

    def __iter__(self):
        l = self.mc.get(self.cache_key)
        return l.keys()
    
    def iterkeys(self):
        l = self.mc.get(self.cache_key)
        return l.keys()

# Stolen from CherryPy
class CyclicTimer(threading._Timer):
    """A thread timer that runs untl it is explicitely stopped.
    """
    def run(self):
        while True:
            self.finished.wait(self.interval)
            if self.finished.isSet():
                break
            self.function(*self.args, **self.kwargs)

class Indexer(object):
    def __init__(self, batch=5):
        """This class is the manager of your index handlers.
        You register your index handlers and then you start the
        indexer daemon which will run at regular interval, polling
        from the global queue members to index.

        If you have your own even mechanism you don't need to
        start the daemon and can simply call `apply_all` to
        perform the same job at will.

        The ``batch`` parameter indicates how many members
        to dequeue during one iteration of indexing.
        """
        self.indexes = {}
        self.index_timer = None
        self.batch = batch

    def register(self, index):
        """Registers a new index handler."""
        self.indexes[index.name] = index

    def unregister(self, index):
        """Removes an index handler"""
        if index.name in self.indexes:
            del self.indexes[index.name]

    def retrieve(self, name):
        """Returns the index handler per its name"""
        return self.indexes[name]

    def shutdown(self):
        """Shutdown all the associated index container""" 
        for name in self.indexes:
            self.indexes[name].container.shutdown()

    def start_daemon(self, interval=60.0):
        """Convenient way to start a daemon timer that will
        process the queue of indexable members.

        Note that if you try to start it before it was stopped
        an error will be raised.
        """
        if self.index_timer is not None:
            raise RuntimeError, "Index daemon already started"
        
        t = CyclicTimer(interval, self.apply_all)
        t.setName('Amplee indexer')
        self.index_timer = t
        self.index_timer.start()

    def stop_daemon(self):
        """Stops the daemon timer"""
        self.shutdown()
        if self.index_timer is not None:
            self.index_timer.cancel()
            self.index_timer.join()
            self.index_timer = None

    def process(self, member):
        """Indexes the provided member"""
        for name in self.indexes:
            self.indexes[name].update(member)
            
    def apply_all(self):
        """Unqueue up to `self.batch` number
        of members from the global queue and apply
        all the index handlers.

        A handler only needs to be a class that implements
        a method named `update` and taking the `member` as its
        unique parameter.
        """
        for unused in xrange(0, self.batch):
            try:
                member = _queue.get_nowait()
            except Empty:
                return
            for name in self.indexes:
                self.indexes[name].update(member)

    def to_dict(self, result):
        """Returns teh given Set as a dictionnary of the form:

        {collection_name: [member_ids]}
        """
        members = {}
        for (collection_name, member_id) in result:
            if collection_name not in members:
                members[collection_name] = []
            members[collection_name].append(member_id)

        return members
    
class BaseIndex(object):
    def __init__(self, name, container=None):
        """
        Base class of your index handler.

        Built-in index handlers will always inherit from
        this class but your own handlers don't need to
        as long as they implement a method named `update`
        and taking the `member` as its unique parameter.

        The ``name`` parameter is an internal identifier.

        The ``container`` is an instance of object
        implementing the dictionnary interface.
        """
        self.name = name
        self.container = container

    def update(self, member):
        raise NotImplemented

    def store(self, key, value):
        """Stores a value within the index container

        The ``key`` and ``value`` can be any object
        as long as the container can cope with it.
        """
        if not key in self.container:
            self.container[key] = Set([value])
        else:
            s = self.container[key]
            s.add(value)
            self.container[key] = s

    def load(self, key):
        """Returns the value associated with the key as a set of one
        element."""
        if key in self.container:
            return self.container[key]

        return Set()

    def iterindex(self, func):
        """Iterates through the keys of the
        container and apply for each one the provided ``func``
        with the key and the data associated.

        If you want to stop the iteration you can
        simply raise StopIteration from the function.

        The method returns a list of results provided
        by the function calls.
        """
        container = None
        result = Set()
        try:
            for key in self.container.iterkeys():
                value = func(key, self.container[key])
                if isinstance(value, Set):
                    result.update(value)
                elif isinstance(value, list):
                    result.update(Set(value))
                elif value is not None:
                    result.add(value)
        except StopIteration:
            pass

        return result

    def keys(self):
        """Return all the existing keys in the container as a list.
        Not all containers support that features and should therefore
        be used carefully.
        """
        return [key for key in self.container.iterkeys()]

class DateIndex(BaseIndex):
    def __init__(self, name, target, container=None, granularity=None):
        """Base class for date based inde handlers.

        The ``target`` is the name of the element within the atom entry
        to search for.

        The ``granularity`` is one of the classmethods of this class
        that indicates what will be the granularity of the key of the index.

        The lower the finer but also the more keys you will have. It
        defaults to ``DateIndex.hour``.
        """
        BaseIndex.__init__(self, name, container)
        if not granularity:
            granularity = DateIndex.hour
        self.granularity = granularity
        self.target = target
        
    def update(self, member):
        entry = member.atom
        published = entry.get_child(self.target, entry.xml_ns)
        if published:
            pub = self.granularity(parse_isodate(published.xml_text))
            self.store(str(pub),  (member.collection.name_or_id, member.member_id))
            
    def between(self, start, end):
        """Returns a dictionnary of the forum {collection_name: [member_id]}
        which have been indexed between the two provided dates.
        """
        def _between(key, data):
            dt = parse_isodate(key)
            if start <= dt <= end:
                return data

        return self.iterindex(_between)

    def year(cls, dt):
        return datetime(dt.year, 1, 1, 0, 0)
    year = classmethod(year)
    
    def month(cls, dt):
        return datetime(dt.year, dt.month, 1, 0, 0)
    month = classmethod(month)
    
    def day(cls, dt):
        return datetime(dt.year, dt.month, dt.day, 0, 0)
    year = classmethod(day)
    
    def hour(cls, dt):
        return datetime(dt.year, dt.month, dt.day, dt.hour, 0)
    hour = classmethod(hour)
    
    def minute(cls, dt):
        return datetime(dt.year, dt.month, dt.day, dt.hour, dt.minute)
    minute = classmethod(minute)
            
class PublishedIndex(DateIndex):
    def __init__(self, name, container=None, granularity=None):
        DateIndex.__init__(self, name, 'published', container, granularity)
    
class UpdatedIndex(DateIndex):
    def __init__(self, name, container=None, granularity=None):
        DateIndex.__init__(self, name, 'updated', container, granularity)
    
class AuthorIndex(BaseIndex):
    def __init__(self, name, container=None, index_name=True, index_uri=False, index_email=False):
        """Simple author indexing.

        This class will create a sha hash value of the concatenation of
        the name, uri and email if the three have been enabled.
        """
        BaseIndex.__init__(self, name, container)
        self.index_name = index_name
        self.index_uri = index_uri
        self.index_email = index_email

    def update(self, member):
        entry = member.atom
        authors = entry.get_children('author', entry.xml_ns)
        if authors:
            for author in authors:
                name = uri = email = ''
                if self.index_name:
                    name = author.get_child('name', author.xml_ns)
                if self.index_uri:
                    uri = author.get_child('uri', author.xml_ns)
                if self.index_email:
                    email = author.get_child('email', author.xml_ns)

                hashed_key = sha.new('%s%s%s' % (name, uri, email)).hexdigest()
                self.store(hashed_key,  (member.collection.name_or_id, member.member_id))
            
    def lookup(self, name='', uri='', email=''):
        hashed_key = sha.new('%s%s%s' % (name, uri, email)).hexdigest()
        return self.load(hashed_key)

class CategoryIndex(BaseIndex):
    def __init__(self, name, container=None, index_term=True, index_scheme=False):
        """Simple category indexing.
        """
        BaseIndex.__init__(self, name, container)
        self.index_term = index_term
        self.index_scheme = index_scheme

    def update(self, member):
        entry = member.atom
        categories = entry.get_children('category', entry.xml_ns)
        if categories:
            for category in categories:
                term = scheme = ''
                if self.index_term:
                    term = category.get_attribute('term')
                if self.index_scheme:
                    scheme = category.get_attribute('scheme')

                hashed_key = sha.new('%s%s' % (term, scheme)).hexdigest()
                self.store(hashed_key, (member.collection.name_or_id, member.member_id))
            
    def lookup(self, term='', scheme=''):
        hashed_key = sha.new('%s%s' % (term, scheme)).hexdigest()
        return self.load(hashed_key)

class KeywordIndex(BaseIndex):
    def __init__(self, name, container=None, keywords=None):
        """Simple keyword indexing on the content element
        """
        BaseIndex.__init__(self, name, container)
        self.keywords = keywords
        self._r = re.compile('|'.join(self.keywords))

    def update(self, member):
        entry = member.atom
        content = entry.get_child('content', entry.xml_ns)
        if content:
            content_type = content.get_attribute('type')
            if unicode(content_type) == u'text':
                match = self._r.search(content.xml_text)
                if match is not None:
                    for keyword in self.keywords:
                        self.store(keyword, (member.collection.name_or_id, member.member_id))
            
    def contains(self, keyword):
        return self.load(keyword)

class AtomIDIndexer(BaseIndex):
    def __init__(self, name, container=None):
        """Indexing atom:id values"""
        BaseIndex.__init__(self, name, container)

    def update(self, member):
        entry = member.atom
        entry_id = entry.get_child('id', entry.xml_ns)
        if entry_id:
            key = entry_id.xml_text.encode('utf-8')
            self.store(key, (member.collection.name_or_id, member.member_id))
        
    def lookup(self, entry_id):
        return self.load(entry_id)