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                 for file_name in os.listdir(dir_name):
  48                     file_path = os.path.join(dir_name, file_name)
  49                     file_base, ext = os.path.splitext(file_path)
  50
  51                     # Skip files that are not XML files
  52                     if not ext == '.xml':
  53                         continue
  54
  55                     if verbose > 0:
  56                         print "Parsing '%s'" % file_path
  57                     else:
  58                         sys.stdout.write('.')
  59                         sys.stdout.flush()
  60
  61                     # Import book files
  62                     try:
  63                         book = Book.from_xml_file(file_path, overwrite=force)
  64                         files_imported += 1
  65
  66                         if os.path.isfile(file_base + '.pdf'):
  67                             book.pdf_file.save('%s.pdf' % book.slug, File(file(file_base + '.pdf')))
  68                             if verbose:
  69                                 print "Importing %s.pdf" % file_base
  70                         if os.path.isfile(file_base + '.odt'):
  71                             book.odt_file.save('%s.odt' % book.slug, File(file(file_base + '.odt')))
  72                             if verbose:
  73                                 print "Importing %s.odt" % file_base
  74                         if os.path.isfile(file_base + '.txt'):
  75                             book.txt_file.save('%s.txt' % book.slug, File(file(file_base + '.txt')))
  76                             if verbose:
  77                                 print "Importing %s.txt" % file_base
  78                         if os.path.isfile(os.path.join(dir_name, book.slug + '.mp3')):
  79                             book.mp3_file.save('%s.mp3' % book.slug, File(file(os.path.join(dir_name, book.slug + '.mp3'))))
  80                             if verbose:
  81                                 print "Importing %s.mp3" % book.slug
  82                         if os.path.isfile(os.path.join(dir_name, book.slug + '.ogg')):
  83                             book.ogg_file.save('%s.ogg' % book.slug, File(file(os.path.join(dir_name, book.slug + '.ogg'))))
  84                             if verbose:
  85                                 print "Importing %s.ogg" % book.slug
  86
  87                         book.save()
  88
  89                     except Book.AlreadyExists, msg:
  90                         print self.style.ERROR('%s: Book already imported. Skipping. To overwrite use --force.' %
  91                             file_path)
  92                         files_skipped += 1
  93
  94         # Print results
  95         print
  96         print "Results: %d files imported, %d skipped, %d total." % (
  97             files_imported, files_skipped, files_imported + files_skipped)
  98         print
  99
 100         transaction.commit()
 101         transaction.leave_transaction_management()
 102