pygal/search.py

from django.conf import settings
import fstools
import logging
from .models import Item
import os
from whoosh.fields import Schema, ID, TEXT, KEYWORD, NUMERIC, DATETIME
from whoosh.qparser.dateparse import DateParserPlugin
from whoosh import index, qparser
from pygal.models import TYPE_FOLDER

logger = logging.getLogger("WHOOSH")


SCHEMA = Schema(
    rel_path=ID(unique=True, stored=True),
    # Item
    name=TEXT,
    type=TEXT,
    favourite_of=KEYWORD,
    datetime=DATETIME,
    size=NUMERIC,
    # Tag
    tag=KEYWORD,
    # Image Cache
    exposure_program=TEXT,
    exposure_time=NUMERIC,
    flash=TEXT,
    f_number=NUMERIC,
    focal_length=NUMERIC,
    lon=NUMERIC,
    lat=NUMERIC,
    height=NUMERIC,
    iso=NUMERIC,
    camera=TEXT,
    orientation=NUMERIC,
    width=NUMERIC,
    # Audio Cache
    album=TEXT,
    artist=TEXT,
    bitrate=NUMERIC,
    duration=NUMERIC,
    genre=TEXT,
    title=TEXT,
    track=NUMERIC,
    year=NUMERIC,
    #
    ratio=NUMERIC,
)


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


def load_index():
        if not os.path.exists(settings.WHOOSH_PATH):
            fstools.mkdir(settings.WHOOSH_PATH)
        try:
            ix = index.open_dir(settings.WHOOSH_PATH)
        except index.EmptyIndexError:
            ix = create_index()
        else:
            logger.debug('Search Index opened.')
        return ix


def item_is_supported(item):
    return item.type != TYPE_FOLDER


def add_item(ix, item):
    # Collect data for the item
    #
    data = {
        'rel_path': item.rel_path,
        'name': os.path.splitext(item.rel_path.split('/')[-1])[0],
        'type': item.type,
    }
    favourite_of = item.favourite_of.all()
    if len(favourite_of) > 0:
        data['favourite_of'] = ' '.join([u.username for u in favourite_of])
    tags = item.tag_set.all()
    if len(tags) > 0:
        data['tag'] = ' '.join([t.text for t in tags])
    for key, value in item.cached_item_data.items():
        data[key] = value
    # Write data to the index
    #
    with ix.writer() as w:
        logger.info('Adding document with rel_path=%s to the search index.', data.get('rel_path'))
        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 rel_path=%s from the search index.', item.rel_path)
        w.delete_by_term("rel_path", item.rel_path)


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


def rebuild_index(ix):
    for item in Item.objects.all().exclude(type=TYPE_FOLDER):
        add_item(ix, item)
    return len(Item.objects.all().exclude(type=TYPE_FOLDER))


def search(ix, search_txt):
    qp = qparser.MultifieldParser(['name', 'tag'], ix.schema)
    qp.add_plugin(DateParserPlugin(free=True))
    q = qp.parse(search_txt)
    with ix.searcher() as s:
        results = s.search(q, limit=None)
        rpl = []
        for hit in results:
            rpl.append(hit['rel_path'])
        return Item.objects.filter(rel_path__in=rpl)