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