slightly prettier annotation links (#677)
[wolnelektury.git] / 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 sorted(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 + '.epub'):
71                             book.epub_file.save('%s.epub' % book.slug, File(file(file_base + '.epub')))
72                             if verbose:
73                                 print "Importing %s.epub" % file_base
74                         if os.path.isfile(file_base + '.odt'):
75                             book.odt_file.save('%s.odt' % book.slug, File(file(file_base + '.odt')))
76                             if verbose:
77                                 print "Importing %s.odt" % file_base
78                         if os.path.isfile(file_base + '.txt'):
79                             book.txt_file.save('%s.txt' % book.slug, File(file(file_base + '.txt')))
80                             if verbose:
81                                 print "Importing %s.txt" % file_base
82                         if os.path.isfile(os.path.join(dir_name, book.slug + '.mp3')):
83                             book.mp3_file.save('%s.mp3' % book.slug, File(file(os.path.join(dir_name, book.slug + '.mp3'))))
84                             if verbose:
85                                 print "Importing %s.mp3" % book.slug
86                         if os.path.isfile(os.path.join(dir_name, book.slug + '.ogg')):
87                             book.ogg_file.save('%s.ogg' % book.slug, File(file(os.path.join(dir_name, book.slug + '.ogg'))))
88                             if verbose:
89                                 print "Importing %s.ogg" % book.slug
90
91                         book.save()
92
93                     except Book.AlreadyExists, msg:
94                         print self.style.ERROR('%s: Book already imported. Skipping. To overwrite use --force.' %
95                             file_path)
96                         files_skipped += 1
97
98         # Print results
99         print
100         print "Results: %d files imported, %d skipped, %d total." % (
101             files_imported, files_skipped, files_imported + files_skipped)
102         print
103
104         transaction.commit()
105         transaction.leave_transaction_management()
106