apps/catalogue/management/commands/importbooks.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 import os
   6 import sys
   7 from optparse import make_option
   8
   9 from django.core.management.base import BaseCommand
  10 from django.core.management.color import color_style
  11 from django.core.files import File
  12
  13 from catalogue.models import Book
  14
  15
  16 class Command(BaseCommand):
  17     option_list = BaseCommand.option_list + (
  18         make_option('-q', '--quiet', action='store_false', dest='verbose', default=True,
  19             help='Verbosity level; 0=minimal output, 1=normal output, 2=all output'),
  20         make_option('-f', '--force', action='store_true', dest='force', default=False,
  21             help='Print status messages to stdout')
  22     )
  23     help = 'Imports books from the specified directories.'
  24     args = 'directory [directory ...]'
  25
  26     def handle(self, *directories, **options):
  27         from django.db import transaction
  28
  29         self.style = color_style()
  30
  31         verbose = options.get('verbose')
  32         force = options.get('force')
  33         show_traceback = options.get('traceback', False)
  34
  35         # Start transaction management.
  36         transaction.commit_unless_managed()
  37         transaction.enter_transaction_management()
  38         transaction.managed(True)
  39
  40         files_imported = 0
  41         files_skipped = 0
  42
  43         for dir_name in directories:
  44             if not os.path.isdir(dir_name):
  45                 print self.style.ERROR("%s: Not a directory. Skipping." % dir_name)
  46             else:
  47                 # files queue
  48                 files = sorted(os.listdir(dir_name))
  49                 postponed = {}
  50                 while files:
  51                     file_name = files.pop(0)
  52                     file_path = os.path.join(dir_name, file_name)
  53                     file_base, ext = os.path.splitext(file_path)
  54
  55                     # Skip files that are not XML files
  56                     if not ext == '.xml':
  57                         continue
  58
  59                     if verbose > 0:
  60                         print "Parsing '%s'" % file_path
  61                     else:
  62                         sys.stdout.write('.')
  63                         sys.stdout.flush()
  64
  65                     # Import book files
  66                     try:
  67                         book = Book.from_xml_file(file_path, overwrite=force)
  68                         files_imported += 1
  69
  70                         if os.path.isfile(file_base + '.pdf'):
  71                             book.pdf_file.save('%s.pdf' % book.slug, File(file(file_base + '.pdf')))
  72                             if verbose:
  73                                 print "Importing %s.pdf" % file_base
  74                         if os.path.isfile(file_base + '.epub'):
  75                             book.epub_file.save('%s.epub' % book.slug, File(file(file_base + '.epub')))
  76                             if verbose:
  77                                 print "Importing %s.epub" % file_base
  78                         if os.path.isfile(file_base + '.odt'):
  79                             book.odt_file.save('%s.odt' % book.slug, File(file(file_base + '.odt')))
  80                             if verbose:
  81                                 print "Importing %s.odt" % file_base
  82                         if os.path.isfile(file_base + '.txt'):
  83                             book.txt_file.save('%s.txt' % book.slug, File(file(file_base + '.txt')))
  84                             if verbose:
  85                                 print "Importing %s.txt" % file_base
  86                         if os.path.isfile(os.path.join(dir_name, book.slug + '.mp3')):
  87                             book.mp3_file.save('%s.mp3' % book.slug, File(file(os.path.join(dir_name, book.slug + '.mp3'))))
  88                             if verbose:
  89                                 print "Importing %s.mp3" % book.slug
  90                         if os.path.isfile(os.path.join(dir_name, book.slug + '.ogg')):
  91                             book.ogg_file.save('%s.ogg' % book.slug, File(file(os.path.join(dir_name, book.slug + '.ogg'))))
  92                             if verbose:
  93                                 print "Importing %s.ogg" % book.slug
  94                         if os.path.isfile(os.path.join(dir_name, book.slug + '.daisy.zip')):
  95                             book.ogg_file.save('%s.daisy.zip' % book.slug, File(file(os.path.join(dir_name, book.slug + '.daisy.zip'))))
  96                             if verbose:
  97                                 print "Importing %s.daisy.zip" % book.slug
  98
  99                         book.save()
 100
 101                     except Book.AlreadyExists, msg:
 102                         print self.style.ERROR('%s: Book already imported. Skipping. To overwrite use --force.' %
 103                             file_path)
 104                         files_skipped += 1
 105
 106                     except Book.DoesNotExist, e:
 107                         if file_name not in postponed or postponed[file_name] < files_imported:
 108                             # push it back into the queue, maybe the missing child will show up
 109                             if verbose:
 110                                 print self.style.NOTICE('Waiting for missing children')
 111                             files.append(file_name)
 112                             postponed[file_name] = files_imported
 113                         else:
 114                             # we're in a loop, nothing's being imported - some child is really missing
 115                             raise e
 116
 117         # Print results
 118         print
 119         print "Results: %d files imported, %d skipped, %d total." % (
 120             files_imported, files_skipped, files_imported + files_skipped)
 121         print
 122
 123         transaction.commit()
 124         transaction.leave_transaction_management()
 125