django-watson bulildwatson_for_model

"""Rebuilds the database indices needed by django-watson."""

from __future__ import unicode_literals, print_function

from django.core.management.base import BaseCommand, CommandError
from django.db import transaction
from django.db.models import get_model

from watson.registration import SearchEngine, _bulk_save_search_entries

class Command(BaseCommand):

    help = "Adds django-watson database indices for a specific model."
    
    @transaction.commit_on_success
    def handle(self, *args, **options):
        """Runs the management command."""
        verbosity = int(options.get("verbosity", 1))
        try:
            model_name = args[0]
        except IndexError:
            raise CommandError, 'FIXME!'
        try:
            engine_slug = args[1]
        except IndexError:
            engine_slug = 'default'
            if verbosity >= 3:
                print('Using search engine "default"')

        # get the engine
        try:
            search_engine = [x[1] for x in SearchEngine.get_created_engines() if x[0] == engine_slug][0]
        except IndexError:
            raise CommandError, 'Search Engine "%s" is not registered!' % engine_slug

        try:
            model = get_model(*model_name.split('.'))  # app label, model name
        except TypeError:  # were we given only model name without app_name?
            registered_models = search_engine.get_registered_models()
            models = [x for x in registered_models if x.__name__ == model_name]
            if len(models) > 1:
                raise CommandError, 'Model name "%s" is not unique, cannot continue!' % model_name
            if models:
                model = models[0]
            else:
                model = None
        if model is None or not search_engine.is_registered(model):
            raise CommandError, 'Model "%s" is not registered with django-watson search engine "%s"!' % (model_name, engine_slug)

        def iter_search_entries():
            local_refreshed_model_count = 0
            for obj in model._default_manager.all().iterator():
                for search_entry in search_engine._update_obj_index_iter(obj):
                    yield search_entry
                local_refreshed_model_count += 1
                if verbosity >= 3:
                    print("Refreshed search entry for {model} {obj} in {engine_slug!r} search engine.".format(
                        model = model._meta.verbose_name,
                        obj = obj,
                        engine_slug = engine_slug,
                    ))
            if verbosity == 2:
                print("Refreshed {local_refreshed_model_count} {model} search entry(s) in {engine_slug!r} search engine.".format(
                    model = model._meta.verbose_name,
                    local_refreshed_model_count = local_refreshed_model_count,
                    engine_slug = engine_slug,
                ))
        _bulk_save_search_entries(iter_search_entries())