turniej elektrybałtów!
[wolnelektury.git] / apps / catalogue / management / commands / pack.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 re
6 import sys
7 from cPickle import load, dump
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 import zipfile
13
14 from catalogue.models import Book, Tag
15
16
17 class Command(BaseCommand):
18     option_list = BaseCommand.option_list + (
19         make_option('-t', '--tags', dest='tags', metavar='SLUG,...',
20             help='Use only books tagged with this tags'),
21         make_option('-i', '--include', dest='include', metavar='SLUG,...',
22             help='Include specific books by slug'),
23         make_option('-e', '--exclude', dest='exclude', metavar='SLUG,...',
24             help='Exclude specific books by slug')
25     )
26     help = 'Prepare data for Lesmianator.'
27     ftypes = ['xml', 'txt', 'html', 'epub', 'pdf']
28     args = '[%s] output_path.zip' % '|'.join(ftypes)
29
30     def handle(self, ftype, path, **options):
31         self.style = color_style()
32         verbose = int(options.get('verbosity'))
33         tags = options.get('tags')
34         include = options.get('include')
35         exclude = options.get('exclude')
36
37         if ftype in self.ftypes:
38             field = "%s_file" % ftype
39         else:
40             print self.style.ERROR('Unknown file type.')
41             return
42
43         books = []
44
45         if include:
46             books += list(Book.objects.filter(slug__in=include.split(',')).only('slug', field))
47
48         if tags:
49             books += list(Book.tagged.with_all(Tag.objects.filter(slug__in=tags.split(','))).only('slug', field))
50         elif not include:
51             books = list(Book.objects.all().only('slug', field))
52
53         if exclude:
54             books = [book for book in books if book.slug not in exclude.split(',')]
55
56         archive = zipfile.ZipFile(path, 'w')
57
58         processed = skipped = 0
59         for book in books:
60             if verbose >= 2:
61                 print 'Parsing', book.slug
62             content = getattr(book, field)
63             if not content:
64                 if verbose >= 1:
65                     print self.style.NOTICE('%s has no %s file' % (book.slug, ftype))
66                 skipped += 1
67                 continue
68             archive.write(content.path, str('%s.%s' % (book.slug, ftype)))
69             processed += 1
70         archive.close()
71
72         if not processed:
73             if skipped:
74                 print self.style.ERROR("No books with %s files found" % ftype)
75             else:
76                 print self.style.ERROR("No books found")
77             return
78
79         if verbose >= 1:
80             print "%d processed, %d skipped" % (processed, skipped)
81             print "Results written to %s" % path