125 строки
3.2 KiB
Python
125 строки
3.2 KiB
Python
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)
|