ecd3fcc97d7ce41ba3232673d1dc442a7e8c5769
[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 import time
8 from optparse import make_option
9
10 from django.core.management.base import BaseCommand
11 from django.core.management.color import color_style
12 from django.core.files import File
13
14 from catalogue.models import Book
15
16
17 class Command(BaseCommand):
18     option_list = BaseCommand.option_list + (
19         make_option('-q', '--quiet', action='store_false', dest='verbose', default=True,
20             help='Verbosity level; 0=minimal output, 1=normal output, 2=all output'),
21         make_option('-f', '--force', action='store_true', dest='force', default=False,
22             help='Print status messages to stdout'),
23         make_option('-E', '--no-build-epub', action='store_false', dest='build_epub', default=True,
24             help='Don\'t build EPUB file'),
25         make_option('-M', '--no-build-mobi', action='store_false', dest='build_mobi', default=True,
26             help='Don\'t build MOBI file'),
27         make_option('-T', '--no-build-txt', action='store_false', dest='build_txt', default=True,
28             help='Don\'t build TXT file'),
29         make_option('-P', '--no-build-pdf', action='store_false', dest='build_pdf', default=True,
30             help='Don\'t build PDF file'),
31         make_option('-w', '--wait-until', dest='wait_until', metavar='TIME',
32             help='Wait until specified time (Y-M-D h:m:s)'),
33     )
34     help = 'Imports books from the specified directories.'
35     args = 'directory [directory ...]'
36
37     def handle(self, *directories, **options):
38         from django.db import transaction
39
40         self.style = color_style()
41
42         verbose = options.get('verbose')
43         force = options.get('force')
44         show_traceback = options.get('traceback', False)
45
46         wait_until = None
47         if options.get('wait_until'):
48             wait_until = time.mktime(time.strptime(options.get('wait_until'), '%Y-%m-%d %H:%M:%S'))
49             if verbose > 0:
50                 print "Will wait until %s; it's %f seconds from now" % (
51                     time.strftime('%Y-%m-%d %H:%M:%S', 
52                     time.localtime(wait_until)), wait_until - time.time())
53
54         # Start transaction management.
55         transaction.commit_unless_managed()
56         transaction.enter_transaction_management()
57         transaction.managed(True)
58
59         files_imported = 0
60         files_skipped = 0
61
62         for dir_name in directories:
63             if not os.path.isdir(dir_name):
64                 print self.style.ERROR("%s: Not a directory. Skipping." % dir_name)
65             else:
66                 # files queue
67                 files = sorted(os.listdir(dir_name))
68                 postponed = {}
69                 while files:
70                     file_name = files.pop(0)
71                     file_path = os.path.join(dir_name, file_name)
72                     file_base, ext = os.path.splitext(file_path)
73
74                     # Skip files that are not XML files
75                     if not ext == '.xml':
76                         continue
77
78                     if verbose > 0:
79                         print "Parsing '%s'" % file_path
80                     else:
81                         sys.stdout.write('.')
82                         sys.stdout.flush()
83
84                     # Import book files
85                     try:
86                         book = Book.from_xml_file(file_path, overwrite=force, 
87                                                   build_epub=options.get('build_epub'),
88                                                   build_txt=options.get('build_txt'),
89                                                   build_pdf=options.get('build_pdf'),
90                                                   build_mobi=options.get('build_mobi'))
91                         files_imported += 1
92
93                         if os.path.isfile(file_base + '.pdf'):
94                             book.pdf_file.save('%s.pdf' % book.slug, File(file(file_base + '.pdf')))
95                             if verbose:
96                                 print "Importing %s.pdf" % file_base
97                         if os.path.isfile(file_base + '.mobi'):
98                             book.mobi_file.save('%s.mobi' % book.slug, File(file(file_base + '.mobi')))
99                             if verbose:
100                                 print "Importing %s.mobi" % file_base
101                         if os.path.isfile(file_base + '.epub'):
102                             book.epub_file.save('%s.epub' % book.slug, File(file(file_base + '.epub')))
103                             if verbose:
104                                 print "Importing %s.epub" % file_base
105                         if os.path.isfile(file_base + '.txt'):
106                             book.txt_file.save('%s.txt' % book.slug, File(file(file_base + '.txt')))
107                             if verbose:
108                                 print "Importing %s.txt" % file_base
109
110                         book.save()
111
112                     except Book.AlreadyExists, msg:
113                         print self.style.ERROR('%s: Book already imported. Skipping. To overwrite use --force.' %
114                             file_path)
115                         files_skipped += 1
116
117                     except Book.DoesNotExist, e:
118                         if file_name not in postponed or postponed[file_name] < files_imported:
119                             # push it back into the queue, maybe the missing child will show up
120                             if verbose:
121                                 print self.style.NOTICE('Waiting for missing children')
122                             files.append(file_name)
123                             postponed[file_name] = files_imported
124                         else:
125                             # we're in a loop, nothing's being imported - some child is really missing
126                             raise e
127
128         # Print results
129         print
130         print "Results: %d files imported, %d skipped, %d total." % (
131             files_imported, files_skipped, files_imported + files_skipped)
132         print
133
134         if wait_until:
135             print 'Waiting...'
136             try:
137                 time.sleep(wait_until - time.time())
138             except IOError:
139                 print "it's already too late"
140
141         transaction.commit()
142         transaction.leave_transaction_management()
143