Housekeeping.
[wolnelektury.git] / src / catalogue / management / commands / importbooks.py
index 338fea6..b8a9aa7 100644 (file)
@@ -1,43 +1,41 @@
-# -*- coding: utf-8 -*-
 # This file is part of Wolnelektury, licensed under GNU Affero GPLv3 or later.
 # Copyright © Fundacja Nowoczesna Polska. See NOTICE for more information.
 #
 import os
 import sys
-import time
-from optparse import make_option
 from django.conf import settings
 from django.core.management.base import BaseCommand
 from django.core.management.color import color_style
 from django.core.files import File
-from catalogue.utils import trim_query_log
+from django.db import transaction
 from librarian.picture import ImageStore
-from wolnelektury.management.profile import profile
 
 from catalogue.models import Book
 from picture.models import Picture
-
 from search.index import Index
 
 
 class Command(BaseCommand):
-    option_list = BaseCommand.option_list + (
-        make_option('-q', '--quiet', action='store_false', dest='verbose', default=True,
-            help='Verbosity level; 0=minimal output, 1=normal output, 2=all output'),
-        make_option('-f', '--force', action='store_true', dest='force', default=False,
-            help='Overwrite works already in the catalogue'),
-        make_option('-D', '--dont-build', dest='dont_build',
-            metavar="FORMAT,...",
-            help="Skip building specified formats"),
-        make_option('-S', '--no-search-index', action='store_false', dest='search_index', default=True,
-            help='Skip indexing imported works for search'),
-        make_option('-w', '--wait-until', dest='wait_until', metavar='TIME',
-            help='Wait until specified time (Y-M-D h:m:s)'),
-        make_option('-p', '--picture', action='store_true', dest='import_picture', default=False,
-            help='Import pictures'),
-    )
     help = 'Imports books from the specified directories.'
-    args = 'directory [directory ...]'
+
+    def add_arguments(self, parser):
+        parser.add_argument(
+                '-q', '--quiet', action='store_false', dest='verbose', default=True,
+                help='Verbosity level; 0=minimal output, 1=normal output, 2=all output')
+        parser.add_argument(
+                '-f', '--force', action='store_true', dest='force',
+                default=False, help='Overwrite works already in the catalogue')
+        parser.add_argument(
+                '-D', '--dont-build', dest='dont_build', metavar="FORMAT,...",
+                help="Skip building specified formats")
+        parser.add_argument(
+                '-S', '--no-search-index', action='store_false',
+                dest='search_index', default=True,
+                help='Skip indexing imported works for search')
+        parser.add_argument(
+                '-p', '--picture', action='store_true', dest='import_picture',
+                default=False, help='Import pictures')
+        parser.add_argument('directory', nargs='+')
 
     def import_book(self, file_path, options):
         verbose = options.get('verbose')
@@ -57,65 +55,48 @@ class Command(BaseCommand):
                     save=False
                     )
                 if verbose:
-                    print "Importing %s.%s" % (file_base, ebook_format)
+                    print("Importing %s.%s" % (file_base, ebook_format))
         book.save()
 
     def import_picture(self, file_path, options, continue_on_error=True):
         try:
             image_store = ImageStore(os.path.dirname(file_path))
             picture = Picture.from_xml_file(file_path, image_store=image_store, overwrite=options.get('force'))
-        except Exception, ex:
+        except Exception as ex:
             if continue_on_error:
-                print "%s: %s" % (file_path, ex)
+                print("%s: %s" % (file_path, ex))
                 return
             else:
                 raise ex
         return picture
 
-    #    @profile
-    def handle(self, *directories, **options):
-        from django.db import transaction
-
+    @transaction.atomic
+    def handle(self, **options):
         self.style = color_style()
 
         verbose = options.get('verbose')
         import_picture = options.get('import_picture')
 
-        wait_until = None
-        if options.get('wait_until'):
-            wait_until = time.mktime(time.strptime(options.get('wait_until'), '%Y-%m-%d %H:%M:%S'))
-            if verbose > 0:
-                print "Will wait until %s; it's %f seconds from now" % (
-                    time.strftime('%Y-%m-%d %H:%M:%S',
-                    time.localtime(wait_until)), wait_until - time.time())
-
-        index = None
         if options.get('search_index') and not settings.NO_SEARCH_INDEX:
             index = Index()
             try:
                 index.index_tags()
                 index.index.commit()
-            except Exception, e:
+            except Exception as e:
                 index.index.rollback()
                 raise e
 
-        # Start transaction management.
-        transaction.commit_unless_managed()
-        transaction.enter_transaction_management()
-        transaction.managed(True)
-
         files_imported = 0
         files_skipped = 0
 
-        for dir_name in directories:
+        for dir_name in options['directory']:
             if not os.path.isdir(dir_name):
-                print self.style.ERROR("%s: Not a directory. Skipping." % dir_name)
+                print(self.style.ERROR("%s: Not a directory. Skipping." % dir_name))
             else:
                 # files queue
                 files = sorted(os.listdir(dir_name))
                 postponed = {}
                 while files:
-                    trim_query_log(0)
                     file_name = files.pop(0)
                     file_path = os.path.join(dir_name, file_name)
                     file_base, ext = os.path.splitext(file_path)
@@ -125,7 +106,7 @@ class Command(BaseCommand):
                         continue
 
                     if verbose > 0:
-                        print "Parsing '%s'" % file_path
+                        print("Parsing '%s'" % file_path)
                     else:
                         sys.stdout.write('.')
                         sys.stdout.flush()
@@ -138,18 +119,18 @@ class Command(BaseCommand):
                             self.import_book(file_path, options)
 
                         files_imported += 1
-                        transaction.commit()
 
                     except (Book.AlreadyExists, Picture.AlreadyExists):
-                        print self.style.ERROR('%s: Book or Picture already imported. Skipping. To overwrite use --force.' %
-                            file_path)
+                        print(self.style.ERROR(
+                            '%s: Book or Picture already imported. Skipping. To overwrite use --force.' %
+                            file_path))
                         files_skipped += 1
 
-                    except Book.DoesNotExist, e:
+                    except Book.DoesNotExist as e:
                         if file_name not in postponed or postponed[file_name] < files_imported:
                             # push it back into the queue, maybe the missing child will show up
                             if verbose:
-                                print self.style.NOTICE('Waiting for missing children')
+                                print(self.style.NOTICE('Waiting for missing children'))
                             files.append(file_name)
                             postponed[file_name] = files_imported
                         else:
@@ -157,17 +138,7 @@ class Command(BaseCommand):
                             raise e
 
         # Print results
-        print
-        print "Results: %d files imported, %d skipped, %d total." % (
-            files_imported, files_skipped, files_imported + files_skipped)
-        print
-
-        if wait_until:
-            print 'Waiting...'
-            try:
-                time.sleep(wait_until - time.time())
-            except IOError:
-                print "it's already too late"
-
-        transaction.commit()
-        transaction.leave_transaction_management()
+        print()
+        print("Results: %d files imported, %d skipped, %d total." % (
+            files_imported, files_skipped, files_imported + files_skipped))
+        print()