apps/catalogue/management/commands/importbooks.py

   1 import os
   2 import sys
   3
   4 from django.core.management.base import BaseCommand
   5 from django.core.management.color import color_style
   6 from optparse import make_option
   7
   8 from catalogue.models import Book
   9
  10
  11 class Command(BaseCommand):
  12     option_list = BaseCommand.option_list + (
  13         make_option('-q', '--quiet', action='store_false', dest='verbose', default=True,
  14             help='Verbosity level; 0=minimal output, 1=normal output, 2=all output'),
  15         make_option('-f', '--force', action='store_true', dest='force', default=False,
  16             help='Print status messages to stdout')
  17     )
  18     help = 'Imports books from the specified directories.'
  19     args = 'directory [directory ...]'
  20
  21     def handle(self, *directories, **options):
  22         from django.db import transaction
  23
  24         self.style = color_style()
  25
  26         verbose = options.get('verbose')
  27         force = options.get('force')
  28         show_traceback = options.get('traceback', False)
  29
  30         # Start transaction management.
  31         transaction.commit_unless_managed()
  32         transaction.enter_transaction_management()
  33         transaction.managed(True)
  34
  35         files_imported = 0
  36         files_skipped = 0
  37
  38         for dir_name in directories:
  39             if not os.path.isdir(dir_name):
  40                 print self.style.ERROR("%s: Not a directory. Skipping." % dir_name)
  41                 files_skipped += 1
  42             else:
  43                 for file_name in os.listdir(dir_name):
  44                     file_path = os.path.join(dir_name, file_name)
  45                     file_base, ext = os.path.splitext(file_path)
  46
  47                     # Skip files that are not XML files
  48                     if not ext == '.xml':
  49                         print self.style.NOTICE("%s: Not an XML file. Skipping." % file_path)
  50                         files_skipped += 1
  51                         continue
  52
  53                     if verbose > 0:
  54                         print "Parsing '%s'" % file_path
  55                     else:
  56                         sys.stdout.write('.')
  57                         sys.stdout.flush()
  58
  59                     # Import book files
  60                     try:
  61                         book = Book.from_xml_file(file_path, overwrite=force)
  62                         files_imported += 1
  63
  64                         if os.path.isfile(file_base + '.pdf'):
  65                             book.pdf_file.save('%s.pdf' % book.slug, File(file(file_base + '.pdf')))
  66                             if verbose:
  67                                 print "Importing %s.pdf" % file_base
  68                         if os.path.isfile(file_base + '.odt'):
  69                             book.odt_file.save('%s.odt' % book.slug, File(file(file_base + '.odt')))
  70                             if verbose:
  71                                 print "Importing %s.odt" % file_base
  72                         if os.path.isfile(file_base + '.txt'):
  73                             book.txt_file.save('%s.txt' % book.slug, File(file(file_base + '.txt')))
  74                             if verbose:
  75                                 print "Importing %s.txt" % file_base
  76
  77                         book.save()
  78
  79                     except Book.AlreadyExists, msg:
  80                         print self.style.ERROR('%s: Book already imported. Skipping. To overwrite use --force.' %
  81                             file_path)
  82                         files_skipped += 1
  83
  84         # Print results
  85         print
  86         print "Results: %d files imported, %d skipped, %d total." % (
  87             files_imported, files_skipped, files_imported + files_skipped)
  88         print
  89
  90         transaction.commit()
  91         transaction.leave_transaction_management()
  92