src/catalogue/management/commands/importbooks.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 os
   5 import sys
   6 from django.conf import settings
   7 from django.core.management.base import BaseCommand
   8 from django.core.management.color import color_style
   9 from django.core.files import File
  10 from django.db import transaction
  11 from librarian.picture import ImageStore
  12
  13 from catalogue.models import Book
  14 from picture.models import Picture
  15 from search.index import Index
  16
  17
  18 class Command(BaseCommand):
  19     help = 'Imports books from the specified directories.'
  20
  21     def add_arguments(self, parser):
  22         parser.add_argument(
  23                 '-q', '--quiet', action='store_false', dest='verbose', default=True,
  24                 help='Verbosity level; 0=minimal output, 1=normal output, 2=all output')
  25         parser.add_argument(
  26                 '-f', '--force', action='store_true', dest='force',
  27                 default=False, help='Overwrite works already in the catalogue')
  28         parser.add_argument(
  29                 '-D', '--dont-build', dest='dont_build', metavar="FORMAT,...",
  30                 help="Skip building specified formats")
  31         parser.add_argument(
  32                 '-S', '--no-search-index', action='store_false',
  33                 dest='search_index', default=True,
  34                 help='Skip indexing imported works for search')
  35         parser.add_argument(
  36                 '-F', '--not-findable', action='store_false',
  37                 dest='findable', default=True,
  38                 help='Set book as not findable.')
  39         parser.add_argument(
  40                 '-p', '--picture', action='store_true', dest='import_picture',
  41                 default=False, help='Import pictures')
  42         parser.add_argument('directory', nargs='+')
  43
  44     def import_book(self, file_path, options):
  45         verbose = options.get('verbose')
  46         if options.get('dont_build'):
  47             dont_build = options.get('dont_build').lower().split(',')
  48         else:
  49             dont_build = None
  50         file_base, ext = os.path.splitext(file_path)
  51         book = Book.from_xml_file(file_path, overwrite=options.get('force'),
  52                                   dont_build=dont_build,
  53                                   search_index_tags=False,
  54                                   findable=options.get('findable'),
  55                                   )
  56         for ebook_format in Book.ebook_formats:
  57             if os.path.isfile(file_base + '.' + ebook_format):
  58                 getattr(book, '%s_file' % ebook_format).save(
  59                     '%s.%s' % (book.slug, ebook_format),
  60                     File(file(file_base + '.' + ebook_format)),
  61                     save=False
  62                     )
  63                 if verbose:
  64                     print("Importing %s.%s" % (file_base, ebook_format))
  65         book.save()
  66
  67     def import_picture(self, file_path, options, continue_on_error=True):
  68         try:
  69             image_store = ImageStore(os.path.dirname(file_path))
  70             picture = Picture.from_xml_file(file_path, image_store=image_store, overwrite=options.get('force'))
  71         except Exception as ex:
  72             if continue_on_error:
  73                 print("%s: %s" % (file_path, ex))
  74                 return
  75             else:
  76                 raise ex
  77         return picture
  78
  79     @transaction.atomic
  80     def handle(self, **options):
  81         self.style = color_style()
  82
  83         verbose = options.get('verbose')
  84         import_picture = options.get('import_picture')
  85
  86         if options.get('search_index') and not settings.NO_SEARCH_INDEX:
  87             index = Index()
  88             try:
  89                 index.index_tags()
  90                 index.index.commit()
  91             except Exception as e:
  92                 index.index.rollback()
  93                 raise e
  94
  95         files_imported = 0
  96         files_skipped = 0
  97
  98         for dir_name in options['directory']:
  99             if not os.path.isdir(dir_name):
 100                 print(self.style.ERROR("%s: Not a directory. Skipping." % dir_name))
 101             else:
 102                 # files queue
 103                 files = sorted(os.listdir(dir_name))
 104                 postponed = {}
 105                 while files:
 106                     file_name = files.pop(0)
 107                     file_path = os.path.join(dir_name, file_name)
 108                     file_base, ext = os.path.splitext(file_path)
 109
 110                     # Skip files that are not XML files
 111                     if not ext == '.xml':
 112                         continue
 113
 114                     if verbose > 0:
 115                         print("Parsing '%s'" % file_path)
 116                     else:
 117                         sys.stdout.write('.')
 118                         sys.stdout.flush()
 119
 120                     # Import book files
 121                     try:
 122                         if import_picture:
 123                             self.import_picture(file_path, options)
 124                         else:
 125                             self.import_book(file_path, options)
 126
 127                         files_imported += 1
 128
 129                     except (Book.AlreadyExists, Picture.AlreadyExists):
 130                         print(self.style.ERROR(
 131                             '%s: Book or Picture already imported. Skipping. To overwrite use --force.' %
 132                             file_path))
 133                         files_skipped += 1
 134
 135                     except Book.DoesNotExist as e:
 136                         if file_name not in postponed or postponed[file_name] < files_imported:
 137                             # push it back into the queue, maybe the missing child will show up
 138                             if verbose:
 139                                 print(self.style.NOTICE('Waiting for missing children'))
 140                             files.append(file_name)
 141                             postponed[file_name] = files_imported
 142                         else:
 143                             # we're in a loop, nothing's being imported - some child is really missing
 144                             raise e
 145
 146         # Print results
 147         print()
 148         print("Results: %d files imported, %d skipped, %d total." % (
 149             files_imported, files_skipped, files_imported + files_skipped))
 150         print()