src/search/management/commands/reindex_pictures.py

   1 # This file is part of Wolnelektury, licensed under GNU Affero GPLv3 or later.
   2 # Copyright © Fundacja Nowoczesna Polska. See NOTICE for more information.
   3 #
   4 import sys
   5 import traceback
   6
   7 from django.core.management.base import BaseCommand
   8
   9
  10 def query_yes_no(question, default="yes"):
  11     """Ask a yes/no question via raw_input() and return their answer.
  12
  13     "question" is a string that is presented to the user.
  14     "default" is the presumed answer if the user just hits <Enter>.
  15         It must be "yes" (the default), "no" or None (meaning
  16         an answer is required of the user).
  17
  18     The "answer" return value is one of "yes" or "no".
  19     """
  20     valid = {"yes": True, "y": True, "ye": True,
  21              "no": False, "n": False}
  22     if default is None:
  23         prompt = " [y/n] "
  24     elif default == "yes":
  25         prompt = " [Y/n] "
  26     elif default == "no":
  27         prompt = " [y/N] "
  28     else:
  29         raise ValueError("invalid default answer: '%s'" % default)
  30
  31     while True:
  32         sys.stdout.write(question + prompt)
  33         choice = raw_input().lower()
  34         if default is not None and choice == '':
  35             return valid[default]
  36         elif choice in valid:
  37             return valid[choice]
  38         else:
  39             sys.stdout.write("Please respond with 'yes' or 'no' (or 'y' or 'n').\n")
  40
  41
  42 class Command(BaseCommand):
  43     help = 'Reindex pictures.'
  44
  45     def add_arguments(self, parser):
  46         self.add_argument(
  47                 '-n', '--picture-id', action='store_true', dest='picture_id',
  48                 default=False, help='picture id instead of slugs')
  49         self.add_argument('slug/id', nargs='*', metavar='slug/id')
  50
  51     def handle(self, **opts):
  52         from picture.models import Picture
  53         from search.index import Index
  54         idx = Index()
  55
  56         if opts['args']:
  57             pictures = []
  58             for a in opts['args']:
  59                 if opts['picture_id']:
  60                     pictures += Picture.objects.filter(id=int(a)).all()
  61                 else:
  62                     pictures += Picture.objects.filter(slug=a).all()
  63         else:
  64             pictures = list(Picture.objects.order_by('slug'))
  65         while pictures:
  66             try:
  67                 p = pictures[0]
  68                 print(p.slug)
  69                 idx.index_picture(p)
  70                 idx.index.commit()
  71                 pictures.pop(0)
  72             except:
  73                 traceback.print_exc()
  74                 try:
  75                     # we might not be able to rollback
  76                     idx.index.rollback()
  77                 except:
  78                     pass
  79                 retry = query_yes_no("Retry?")
  80                 if not retry:
  81                     break