src/catalogue/helpers.py

   1 # -*- coding: utf-8 -*-
   2 # This file is part of Wolnelektury, licensed under GNU Affero GPLv3 or later.
   3 # Copyright © Fundacja Nowoczesna Polska. See NOTICE for more information.
   4 #
   5 from django.conf import settings
   6 from django.contrib.contenttypes.models import ContentType
   7 from django.core.cache import cache
   8
   9 from .models import Tag, Book
  10 from os.path import getmtime
  11 import cPickle
  12 from collections import defaultdict
  13
  14
  15 BOOK_CATEGORIES = ('author', 'epoch', 'genre', 'kind')
  16
  17 _COUNTERS = None
  18 _COUNTER_TIME = None
  19
  20
  21 def get_top_level_related_tags(tags, categories=None):
  22     """
  23     Finds tags related to given tags through books, and counts their usage.
  24
  25     Takes ancestry into account: if a tag is applied to a book, its
  26     usage on the book's descendants is ignored.
  27     """
  28     global _COUNTERS, _COUNTER_TIME
  29     # First, check that we have a valid and recent version of the counters.
  30     if getmtime(settings.CATALOGUE_COUNTERS_FILE) > _COUNTER_TIME:
  31         with open(settings.CATALOGUE_COUNTERS_FILE) as f:
  32             _COUNTERS = cPickle.load(f)
  33
  34     tagids = tuple(sorted(t.pk for t in tags))
  35     try:
  36         related_ids = _COUNTERS['next'][tagids]
  37     except KeyError:
  38         return
  39
  40     related = Tag.objects.filter(pk__in=related_ids)
  41
  42     # TODO: do we really need that?
  43     if categories is not None:
  44         related = related.filter(category__in=categories)
  45
  46     for tag in related:
  47         tag.count = _COUNTERS['count'][tuple(sorted(tagids + (tag.pk,)))]
  48         yield tag
  49
  50
  51 def update_counters():
  52     def combinations(things):
  53         if len(things):
  54             for c in combinations(things[1:]):
  55                 yield c
  56                 yield (things[0],) + c
  57         else:
  58             yield ()
  59
  60     def count_for_book(book, count_by_combination=None, parent_combinations=None):
  61         if not parent_combinations:
  62             parent_combinations = set()
  63         tags = sorted(tuple(t.pk for t in book.tags.filter(category__in=('author', 'genre', 'epoch', 'kind'))))
  64         combs = list(combinations(tags))
  65         for c in combs:
  66             if c not in parent_combinations:
  67                 count_by_combination[c] += 1
  68         combs_for_child = set(list(parent_combinations) + combs)
  69         for child in book.children.all():
  70             count_for_book(child, count_by_combination, combs_for_child)
  71
  72     count_by_combination = defaultdict(lambda: 0)
  73     for b in Book.objects.filter(parent=None):
  74         count_for_book(b, count_by_combination)
  75
  76     next_combinations = defaultdict(set)
  77     # Now build an index of all combinations.
  78     for c in count_by_combination.keys():
  79         if not c:
  80             continue
  81         for n in c:
  82             rest = tuple(x for x in c if x != n)
  83             next_combinations[rest].add(n)
  84
  85     counters = {
  86         "count": dict(count_by_combination),
  87         "next": dict(next_combinations),
  88     }
  89
  90     with open(settings.CATALOGUE_COUNTERS_FILE, 'w') as f:
  91         cPickle.dump(counters, f)
  92
  93
  94 def get_audiobook_tags():
  95     audiobook_tag_ids = cache.get('audiobook_tags')
  96     if audiobook_tag_ids is None:
  97         books_with_audiobook = Book.objects.filter(media__type__in=('mp3', 'ogg'))\
  98             .distinct().values_list('pk', flat=True)
  99         audiobook_tag_ids = Tag.objects.filter(
 100             items__content_type=ContentType.objects.get_for_model(Book),
 101             items__object_id__in=list(books_with_audiobook)).distinct().values_list('pk', flat=True)
 102         audiobook_tag_ids = list(audiobook_tag_ids)
 103         cache.set('audiobook_tags', audiobook_tag_ids)
 104     return audiobook_tag_ids