src/catalogue/management/commands/pack.py

   1 # This file is part of Wolne Lektury, licensed under GNU Affero GPLv3 or later.
   2 # Copyright © Fundacja Wolne Lektury. 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(
  24                 '--top-level', dest='top_level', action='store_true')
  25         parser.add_argument('ftype', metavar='|'.join(Book.formats))
  26         parser.add_argument('path', metavar='output_path.zip')
  27
  28     def handle(self, **options):
  29         self.style = color_style()
  30         ftype = options['ftype']
  31         path = options['path']
  32         verbose = int(options.get('verbosity'))
  33         tags = options.get('tags')
  34         include = options.get('include')
  35         exclude = options.get('exclude')
  36         top_level = options.get('top_level')
  37
  38         if ftype in Book.formats:
  39             field = "%s_file" % ftype
  40         else:
  41             print(self.style.ERROR('Unknown file type.'))
  42             return
  43
  44         books = []
  45
  46         if include:
  47             books += list(Book.objects.filter(slug__in=include.split(',')).only('slug', field))
  48
  49         if tags:
  50             books += list(Book.tagged.with_all(Tag.objects.filter(slug__in=tags.split(','))).only('slug', field))
  51         elif not include:
  52             books = Book.objects.all()
  53             if top_level:
  54                 books = books.filter(parent=None)
  55             books = list(books.only('slug', field))
  56
  57         if exclude:
  58             books = [book for book in books if book.slug not in exclude.split(',')]
  59
  60         archive = zipfile.ZipFile(path, 'w', zipfile.ZIP_DEFLATED)
  61
  62         processed = skipped = 0
  63         for book in books:
  64             if verbose >= 2:
  65                 print('Parsing', book.slug)
  66             content = getattr(book, field)
  67             if not content:
  68                 if verbose >= 1:
  69                     print(self.style.NOTICE('%s has no %s file' % (book.slug, ftype)))
  70                 skipped += 1
  71                 continue
  72             archive.write(content.path, str('%s.%s' % (book.slug, ftype)))
  73             processed += 1
  74         archive.close()
  75
  76         if not processed:
  77             if skipped:
  78                 print(self.style.ERROR("No books with %s files found" % ftype))
  79             else:
  80                 print(self.style.ERROR("No books found"))
  81             return
  82
  83         if verbose >= 1:
  84             print("%d processed, %d skipped" % (processed, skipped))
  85             print("Results written to %s" % path)