patt/search.py

import datetime
from django.conf import settings
import fstools
import logging
from .models import Task, TASKSTATE_OPEN, TASKSTATE_FINISHED, TASKSTATE_CLOSED, TASKSTATE_CANCELED
import os
import re
from whoosh.fields import Schema, ID, TEXT, NUMERIC, DATETIME, BOOLEAN
from whoosh.qparser.dateparse import DateParserPlugin
from whoosh import index, qparser

logger = logging.getLogger("WHOOSH")

SEARCH_MY_OPEN_TASKS = 1
SEARCH_LOST_SOULS = 2


def common_searches(request):
    cs = {}
    if request.user.is_authenticated:
        cs[SEARCH_MY_OPEN_TASKS] = ('My Open Tasks', mk_search_pattern(user_ids=[request.user.id]))
    cs[SEARCH_LOST_SOULS] = ('Lost Souls (no user)', 'assigned_user_missing:1')
    return cs


# INDEX_STATES
#
INDEX_STATES = {
    TASKSTATE_OPEN: 'Open',
    TASKSTATE_FINISHED: 'Finished',
    TASKSTATE_CLOSED: 'Closed',
    TASKSTATE_CANCELED: 'Canselled'
}


SCHEMA = Schema(
    id=ID(unique=True, stored=True),
    # Task
    task_id=NUMERIC,
    assigned_user=TEXT,
    assigned_user_missing=BOOLEAN,
    name=TEXT,
    description=TEXT,
    state=TEXT,
    targetdate=DATETIME,
    # Related Project
    project_id=NUMERIC,
    project_name=TEXT,
    project_description=TEXT,
    # Related Comments
    comment=TEXT,
)


def mk_whooshpath_if_needed():
    if not os.path.exists(settings.WHOOSH_PATH):
        fstools.mkdir(settings.WHOOSH_PATH)


def create_index():
    mk_whooshpath_if_needed()
    logger.debug('Search Index created.')
    return index.create_in(settings.WHOOSH_PATH, schema=SCHEMA)


def load_index():
        mk_whooshpath_if_needed()
        try:
            ix = index.open_dir(settings.WHOOSH_PATH)
        except index.EmptyIndexError:
            ix = create_index()
        else:
            logger.debug('Search Index opened.')
        return ix


def add_item(ix, item):
    # Define Standard data
    #
    data = dict(
        id='%d' % item.id,
        # Task
        task_id=item.id,
        name=item.name,
        description=item.description,
        state=INDEX_STATES.get(item.state),
        # Related Project
        project_id=item.project.id,
        project_name=item.project.name,
        project_description=item.project.description,
        # Related Comments
        comment=' '.join([c.comment for c in item.comment_set.all()]),
    )
    # Add Optional data
    #
    if item.assigned_user is not None:
        data['assigned_user'] = item.assigned_user.username
        data['assigned_user_missing'] = False
    else:
        data['assigned_user_missing'] = True
    if item.targetdate is not None:
        data['targetdate'] = datetime.datetime.combine(item.targetdate, datetime.datetime.min.time())
    # Write data to the index
    #
    with ix.writer() as w:
        logger.info('Adding document with id=%d to the search index.', data.get('task_id'))
        w.add_document(**data)
        for key in data:
            logger.debug('  - Adding %s=%s', key, repr(data[key]))


def delete_item(ix, item):
    with ix.writer() as w:
        logger.info('Removing document with id=%d from the search index.', item.id)
        w.delete_by_term("task_id", item.id)


def update_item(ix, item):
    delete_item(ix, item)
    add_item(ix, item)


def rebuild_index(ix):
    for t in Task.objects.all():
        add_item(ix, t)
    return len(Task.objects.all())


def search(ix, search_txt):
    qp = qparser.MultifieldParser(['name', 'description'], ix.schema)
    qp.add_plugin(DateParserPlugin(free=True))
    try:
        q = qp.parse(search_txt)
    except AttributeError:
        return None
    except Exception:
        return None
    with ix.searcher() as s:
        results = s.search(q, limit=None)
        rpl = []
        for hit in results:
            rpl.append(hit['id'])
        return Task.objects.filter(id__in=rpl)


def mk_search_pattern(**kwargs):
    prj_ids = kwargs.get('prj_ids', [])
    user_ids = kwargs.get('user_ids', [])
    states = kwargs.get('states', [INDEX_STATES.get(TASKSTATE_OPEN), INDEX_STATES.get(TASKSTATE_FINISHED)])
    rule_parts = []
    if prj_ids is not None and len(prj_ids) > 0:
        rule_parts.append(' OR '.join(['project_id:%s' % pid for pid in prj_ids]))
    if user_ids is not None and len(user_ids) > 0:
        from django.contrib.auth.models import User
        rule_parts.append(' OR '.join(['assigned_user:%s' % User.objects.get(id=uid).username for uid in user_ids]))
    if states is not None and len(states) > 0:
        rule_parts.append(' OR '.join(['state:%s' % state for state in states]))
    return ' AND '.join('(%s)' % rule for rule in rule_parts)


def get_project_ids_from_search_pattern(search_txt):
    try:
        return re.findall('project_id:(\d+)', search_txt)
    except AttributeError:
        return None
    except TypeError:
        return None