Merge branch 'commands' into 'master'

Remove duplicate argument validations in commands

See merge request !204

arkindex/dataimport/management/commands/import.py

 #!/usr/bin/env python3
-from django.core.management.base import BaseCommand, CommandError
+from django.core.management.base import BaseCommand
 from arkindex_common.ml_tool import MLToolType
-from arkindex.dataimport.models import DataImport, DataImportMode
+from arkindex.project.argparse import DataImportArgument
+from arkindex.dataimport.models import DataImportMode
 from arkindex.dataimport.tasks import extract_pdf_images, populate_volume, setup_ml_analysis, check_images
 from arkindex.dataimport.git import GitFlow
 from django.conf import settings
@@ -23,22 +24,17 @@ class Command(BaseCommand):
     def add_arguments(self, parser):
         parser.add_argument(
             'data_import',
-            type=str,
+            type=DataImportArgument(),
             help='ID of the DataImport to run'
         )
 
-    def handle(self, *args, **options):
+    def handle(self, *args, data_import=None, **options):
         # Use default ML tools here
         ml_tools = (
             (MLToolType.Classifier, settings.ML_DEFAULT_CLASSIFIER),
             (MLToolType.Recognizer, settings.ML_DEFAULT_RECOGNIZER),
         )
 
-        try:
-            data_import = DataImport.objects.get(pk=options['data_import'])
-        except DataImport.DoesNotExist:
-            raise CommandError('Missing DataImport')
-
         # Use shared directory when running in docker
         # Fallback to a temp directory while developing
         task_dir = os.environ.get('PONOS_DATA', tempfile.mkdtemp(suffix='-ponos'))

arkindex/dataimport/management/commands/import_repo.py

-#!/usr/bin/env python3
 from django.core.management.base import BaseCommand, CommandError
-from arkindex.dataimport.models import DataImport, DataImportMode, Repository
+from arkindex.project.argparse import RepositoryArgument
+from arkindex.dataimport.models import DataImport, DataImportMode
 
 
 class Command(BaseCommand):
@@ -8,8 +8,9 @@ class Command(BaseCommand):
 
     def add_arguments(self, parser):
         parser.add_argument(
-            'repository',
-            help='ID of the repository to check on',
+            'repo',
+            type=RepositoryArgument(),
+            help='ID or part of the URL of the repository to check on',
         )
         parser.add_argument(
             '--hash',
@@ -17,14 +18,9 @@ class Command(BaseCommand):
             default=None,
         )
 
-    def handle(self, *args, **options):
-        try:
-            repo = Repository.objects.get(id=options['repository'])
-        except Repository.DoesNotExist:
-            raise CommandError('Repository {} not found'.format(options['repository']))
-
+    def handle(self, *args, repo=None, **options):
         if repo.provider_class is None:
-            raise ValueError("No repository provider found for {}".format(repo.url))
+            raise CommandError("No repository provider found for {}".format(repo.url))
 
         if 'hash' in options and options['hash'] is not None:
             rev, created = repo.provider.get_or_create_revision(repo, options['hash'])
@@ -32,7 +28,7 @@ class Command(BaseCommand):
             rev, created = repo.provider.get_or_create_latest_revision(repo)
 
         if created:
-            print('Created revision {} "{}" on repository {}'.format(rev.hash, rev.message, repo.url))
+            self.stdout.write('Created revision {} "{}" on repository {}'.format(rev.hash, rev.message, repo.url))
 
         di = DataImport.objects.create(
             creator=repo.credentials.user,
@@ -42,5 +38,5 @@ class Command(BaseCommand):
         )
         di.start()
 
-        print('Successfully built DataImport {}'.format(di))
-        print('To test the import manually, run: ./manage.py import {}'.format(di.id))
+        self.stdout.write(self.style.SUCCESS('Successfully built DataImport {}'.format(di)))
+        self.stdout.write('To test the import manually, run: ./manage.py import {}'.format(di.id))

arkindex/documents/management/commands/generate_thumbnails.py

 #!/usr/bin/env python3
 from django.core.management.base import CommandError
-from django.core.exceptions import ValidationError
 from django.conf import settings
 from ponos.management.base import PonosCommand
-from arkindex.documents.models import Corpus, Element, ElementType
+from arkindex.project.argparse import CorpusArgument, ElementArgument
+from arkindex.documents.models import Element, ElementType
 
 
 class Command(PonosCommand):
@@ -13,20 +13,21 @@ class Command(PonosCommand):
 
     def add_arguments(self, parser):
         super().add_arguments(parser)
-        # TODO: Mutually exclusive group
+        parser.add_argument(
+            '--all',
+            help='Create thumbnails for every volume in every corpus',
+            action='store_true',
+            default=False,
+        )
         parser.add_argument(
             '--corpus',
             help='ID or part of the name of the corpus to fetch volumes from',
+            type=CorpusArgument(),
         )
         parser.add_argument(
             '--element',
             help='ID or part of the name of a single element to build a thumbnail for',
-        )
-        parser.add_argument(
-            '--all',
-            help='Create thumbnails for every volume in every corpus',
-            action='store_true',
-            default=False,
+            type=ElementArgument(),
         )
         parser.add_argument(
             '--force',
@@ -35,34 +36,23 @@ class Command(PonosCommand):
             default=False,
         )
 
-    def validate_args(self, **options):
-        if options['all']:
-            if options['corpus'] or options['element']:
+    def validate_args(self, corpus=None, element=None, all=False, force=False, **options):
+        if all:
+            if corpus or element:
                 raise CommandError('--all cannot be used together with --corpus or --element')
-            return {'elements': Element.objects.filter(type=ElementType.Volume), 'force': options['force']}
+            return {'elements': Element.objects.filter(type=ElementType.Volume), 'force': force}
 
-        try:
-            corpus = Corpus.objects.get(pk=options['corpus'])
-        except (Corpus.DoesNotExist, ValidationError):
-            try:
-                corpus = Corpus.objects.get(name__icontains=options['corpus'])
-            except Corpus.DoesNotExist:
-                raise CommandError('Corpus "{}" does not exist'.format(options['corpus']))
-
-        if options['element']:
-            try:
-                elt = corpus.elements.get(pk=options['element'])
-            except (Element.DoesNotExist, ValidationError):
-                try:
-                    elt = corpus.elements.get(name__icontains=options['element'])
-                except Element.DoesNotExist:
-                    raise CommandError('Element "{}" does not exist'.format(options['element']))
+        if not corpus:
+            raise CommandError('--corpus is required when not using --all')
 
-            elts = [elt, ]
+        if element:
+            if not corpus.elements.filter(pk=element.pk).exists():
+                raise CommandError('Element {} is not in corpus {}'.format(element, corpus))
+            elts = [element, ]
         else:
-            elts = Element.objects.filter(corpus=corpus, type=ElementType.Volume)
+            elts = corpus.elements.filter(type=ElementType.Volume)
 
-        return {'elements': elts, 'force': options['force']}
+        return {'elements': elts, 'force': force}
 
     def run(self, elements=[], force=False):
         for element in elements:

arkindex/documents/management/commands/reindex.py

 #!/usr/bin/env python3
-from django.core.management.base import CommandError
-from django.core.exceptions import ValidationError
 from django.conf import settings
 from ponos.management.base import PonosCommand
+from arkindex.project.argparse import ElementArgument
 from arkindex.documents.indexer import Indexer
 from arkindex.documents.models import Element, ElementType, Act, Transcription, Page
 import logging
@@ -87,6 +86,7 @@ class Command(PonosCommand):
         parser.add_argument(
             '--volume',
             help='Restrict reindexing to a specific volume by ID or part of the name',
+            type=ElementArgument(type=ElementType.Volume),
         )
         parser.add_argument(
             '--drop',
@@ -100,26 +100,13 @@ class Command(PonosCommand):
             for k in self.index_methods.keys():
                 options[k] = True
 
-        volume = None
-        if options['volume']:
-            try:
-                volume = Element.objects.get(type=ElementType.Volume, pk=options['volume'])
-            except (Element.DoesNotExist, ValidationError):
-                try:
-                    volume = Element.objects.get(
-                        type=ElementType.Volume,
-                        name__icontains=options['volume'],
-                    )
-                except Element.DoesNotExist:
-                    raise CommandError('Volume "{}" not found'.format(options['volume']))
-
         return {
             'methods': [
                 key
                 for key in self.index_methods.keys()
                 if options.get(key)
             ],
-            'volume': volume,
+            'volume': options['volume'],
             'drop': options.get('drop', False),
         }
 

arkindex/documents/tests/commands/test_generate_thumbnails.py

 from django.core.management import call_command
+from django.core.management.base import CommandError
 from unittest.mock import patch, call
 from arkindex.project.tests import FixtureTestCase
-from arkindex.documents.models import Element, ElementType
+from arkindex.documents.models import Corpus, ElementType
 
 
 class TestGenerateThumbnailsCommand(FixtureTestCase):
@@ -12,9 +13,11 @@ class TestGenerateThumbnailsCommand(FixtureTestCase):
     @classmethod
     def setUpTestData(cls):
         super().setUpTestData()
-        cls.reg = Element.objects.get(type=ElementType.Register, name="Register 1")
-        cls.vol1 = Element.objects.get(type=ElementType.Volume, name="Volume 1")
-        cls.vol2 = Element.objects.get(type=ElementType.Volume, name="Volume 2")
+        cls.reg = cls.corpus.elements.get(type=ElementType.Register, name="Register 1")
+        cls.vol1 = cls.corpus.elements.get(type=ElementType.Volume, name="Volume 1")
+        cls.vol2 = cls.corpus.elements.get(type=ElementType.Volume, name="Volume 2")
+        corpus2 = Corpus.objects.create(name='Other corpus')
+        cls.vol3 = corpus2.elements.create(type=ElementType.Volume, name='Volume 3')
         cls.thumb_patch = patch('arkindex.images.models.Thumbnail.create')
 
     def setUp(self):
@@ -27,11 +30,11 @@ class TestGenerateThumbnailsCommand(FixtureTestCase):
 
     def test_start_corpus(self):
         """
-        Test generate_thumbnails starts a Celery task for each volume of a corpus
+        Test generate_thumbnails runs generation for each volume of a corpus
         """
         call_command(
             'generate_thumbnails',
-            corpus=str(self.corpus.id),
+            corpus=self.corpus,
         )
         self.assertCountEqual(self.thumb_mock.call_args_list, [
             call(self.vol1),
@@ -40,54 +43,71 @@ class TestGenerateThumbnailsCommand(FixtureTestCase):
 
     def test_start_element(self):
         """
-        Test generate_thumbnails starts a Celery task for an element
+        Test generate_thumbnails runs generation for an element
         """
         call_command(
             'generate_thumbnails',
-            corpus=str(self.corpus.id),
-            element=str(self.reg.id),
+            corpus=self.corpus,
+            element=self.reg,
         )
         self.assertCountEqual(self.thumb_mock.call_args_list, [
             call(self.reg),
         ])
 
-    def test_corpus_name(self):
+    @patch('arkindex.documents.models.Element.generate_thumbnail')
+    def test_force(self, gen_mock):
         """
-        Test generate_thumbnails accepts a part of the corpus name instead of an ID
+        Test generate_thumbnails passes the --force argument to tasks
         """
         call_command(
             'generate_thumbnails',
-            corpus='tests',
+            corpus=self.corpus,
+            element=self.reg,
+            force=True,
         )
-        self.assertCountEqual(self.thumb_mock.call_args_list, [
-            call(self.vol1),
-            call(self.vol2),
+        self.assertCountEqual(gen_mock.call_args_list, [
+            call(force=True),
         ])
 
-    def test_element_name(self):
+    def test_all(self):
         """
-        Test generate_thumbnails accepts a part of an element name instead of an ID
+        Test generate_thumbnails picks all volumes when using --all
         """
         call_command(
             'generate_thumbnails',
-            corpus=str(self.corpus.id),
-            element='register 1',
+            all=True,
         )
         self.assertCountEqual(self.thumb_mock.call_args_list, [
-            call(self.reg),
+            call(self.vol1),
+            call(self.vol2),
+            call(self.vol3),
         ])
 
-    @patch('arkindex.documents.models.Element.generate_thumbnail')
-    def test_force(self, gen_mock):
+    def test_all_xor_corpus(self):
         """
-        Test generate_thumbnails passes the --force argument to tasks
+        Test generate_thumbnails does not allow --all and --corpus simultaneously
         """
-        call_command(
-            'generate_thumbnails',
-            corpus=str(self.corpus.id),
-            element='register 1',
-            force=True,
-        )
-        self.assertCountEqual(gen_mock.call_args_list, [
-            call(force=True),
-        ])
+        with self.assertRaisesRegex(CommandError, r'--all.+--corpus'):
+            call_command(
+                'generate_thumbnails',
+                all=True,
+                corpus=self.corpus,
+            )
+
+    def test_corpus_required(self):
+        """
+        Test generate_thumbnails requires either --all or --corpus
+        """
+        with self.assertRaisesRegex(CommandError, '--corpus'):
+            call_command('generate_thumbnails')
+
+    def test_element_in_corpus(self):
+        """
+        Test generate_thumbnails requires --element to be inside --corpus
+        """
+        with self.assertRaisesRegex(CommandError, 'not in corpus'):
+            call_command(
+                'generate_thumbnails',
+                corpus=self.corpus,
+                element=self.vol3,
+            )

arkindex/documents/tests/commands/test_reindex.py

@@ -181,25 +181,13 @@ class TestReindexCommand(FixtureTestCase):
         self.assertEqual(self.indexer_mock().setup.call_count, 0)
         self._assert_all()
 
-    def test_volume_id(self):
+    def test_volume(self):
         """
-        Test the reindex command can restrict indexing to a specific volume by ID
+        Test the reindex command can restrict indexing to a specific volume
         """
         call_command(
             'reindex',
-            volume=str(self.vol.id),
-        )
-        self.assertEqual(self.indexer_mock().drop_index.call_count, 0)
-        self.assertEqual(self.indexer_mock().setup.call_count, 0)
-        self._assert_volume()
-
-    def test_volume_name(self):
-        """
-        Test the reindex command can restrict indexing to a specific volume by name
-        """
-        call_command(
-            'reindex',
-            volume="volume 1",
+            volume=self.vol,
         )
         self.assertEqual(self.indexer_mock().drop_index.call_count, 0)
         self.assertEqual(self.indexer_mock().setup.call_count, 0)

arkindex/project/argparse.py

@@ -2,6 +2,7 @@ from django.db.models import Model
 from django.core.exceptions import ValidationError
 from django.core.management.base import CommandError
 from arkindex.documents.models import Corpus, Element, DataSource
+from arkindex.dataimport.models import DataImport, Repository
 
 
 class ModelArgument(object):
@@ -20,17 +21,22 @@ class ModelArgument(object):
         except (ValidationError, Element.DoesNotExist):
             pass
 
-        text_filter = {'{}__icontains'.format(self.text_search_field): arg}
-        if self.many:
-            return qs.filter(**text_filter)
-
         try:
-            return qs.get(**text_filter)
+            return self.text_search(qs, arg)
         except self.model.DoesNotExist:
             raise CommandError('{} "{}" does not exist'.format(self.model.__name__, arg))
         except self.model.MultipleObjectsReturned:
             raise CommandError('"{}" matches multiple {} instances'.format(arg, self.model.__name__))
 
+    def text_search(self, qs, arg):
+        if not self.text_search_field:
+            raise self.model.DoesNotExist
+        text_filter = {'{}__icontains'.format(self.text_search_field): arg}
+        if self.many:
+            return qs.filter(**text_filter)
+        else:
+            return qs.get(**text_filter)
+
 
 class CorpusArgument(ModelArgument):
     model = Corpus
@@ -43,3 +49,13 @@ class ElementArgument(ModelArgument):
 class DataSourceArgument(ModelArgument):
     model = DataSource
     text_search_field = 'slug'
+
+
+class DataImportArgument(ModelArgument):
+    model = DataImport
+    text_search_field = None
+
+
+class RepositoryArgument(ModelArgument):
+    model = Repository
+    text_search_field = 'url'