General A/B testing.
[wolnelektury.git] / src / catalogue / management / commands / pack.py
1 # This file is part of Wolnelektury, licensed under GNU Affero GPLv3 or later.
2 # Copyright © Fundacja Nowoczesna Polska. See NOTICE for more information.
3 #
4 import zipfile
5 from django.core.management.base import BaseCommand
6 from django.core.management.color import color_style
7 from catalogue.models import Book, Tag
8
9
10 class Command(BaseCommand):
11     help = 'Prepare ZIP package with files of given type.'
12
13     def add_arguments(self, parser):
14         parser.add_argument(
15                 '-t', '--tags', dest='tags', metavar='SLUG,...',
16                 help='Use only books tagged with this tags')
17         parser.add_argument(
18                 '-i', '--include', dest='include', metavar='SLUG,...',
19                 help='Include specific books by slug')
20         parser.add_argument(
21                 '-e', '--exclude', dest='exclude', metavar='SLUG,...',
22                 help='Exclude specific books by slug')
23         parser.add_argument('ftype', metavar='|'.join(Book.formats))
24         parser.add_argument('path', metavar='output_path.zip')
25
26     def handle(self, **options):
27         self.style = color_style()
28         ftype = options['ftype']
29         path = options['path']
30         verbose = int(options.get('verbosity'))
31         tags = options.get('tags')
32         include = options.get('include')
33         exclude = options.get('exclude')
34
35         if ftype in Book.formats:
36             field = "%s_file" % ftype
37         else:
38             print(self.style.ERROR('Unknown file type.'))
39             return
40
41         books = []
42
43         if include:
44             books += list(Book.objects.filter(slug__in=include.split(',')).only('slug', field))
45
46         if tags:
47             books += list(Book.tagged.with_all(Tag.objects.filter(slug__in=tags.split(','))).only('slug', field))
48         elif not include:
49             books = list(Book.objects.all().only('slug', field))
50
51         if exclude:
52             books = [book for book in books if book.slug not in exclude.split(',')]
53
54         archive = zipfile.ZipFile(path, 'w', zipfile.ZIP_DEFLATED)
55
56         processed = skipped = 0
57         for book in books:
58             if verbose >= 2:
59                 print('Parsing', book.slug)
60             content = getattr(book, field)
61             if not content:
62                 if verbose >= 1:
63                     print(self.style.NOTICE('%s has no %s file' % (book.slug, ftype)))
64                 skipped += 1
65                 continue
66             archive.write(content.path, str('%s.%s' % (book.slug, ftype)))
67             processed += 1
68         archive.close()
69
70         if not processed:
71             if skipped:
72                 print(self.style.ERROR("No books with %s files found" % ftype))
73             else:
74                 print(self.style.ERROR("No books found"))
75             return
76
77         if verbose >= 1:
78             print("%d processed, %d skipped" % (processed, skipped))
79             print("Results written to %s" % path)