1from django.apps import apps
2from django.conf import settings
3from django.contrib import admin
4from django.core.management.base import BaseCommand, CommandError
5from reversion.revisions import is_registered
6
7
8class BaseRevisionCommand(BaseCommand):
9
10    def add_arguments(self, parser):
11        super().add_arguments(parser)
12        parser.add_argument(
13            "app_label",
14            metavar="app_label",
15            nargs="*",
16            help="Optional app_label or app_label.model_name list.",
17        )
18        parser.add_argument(
19            "--using",
20            default=None,
21            help="The database to query for revision data.",
22        )
23        parser.add_argument(
24            "--model-db",
25            default=None,
26            help="The database to query for model data.",
27        )
28
29    def get_models(self, options):
30        # Load admin classes.
31        if "django.contrib.admin" in settings.INSTALLED_APPS:
32            admin.autodiscover()
33        # Get options.
34        app_labels = options["app_label"]
35        # Parse model classes.
36        if len(app_labels) == 0:
37            selected_models = apps.get_models()
38        else:
39            selected_models = set()
40            for label in app_labels:
41                if "." in label:
42                    # This is an app.Model specifier.
43                    try:
44                        model = apps.get_model(label)
45                    except LookupError:
46                        raise CommandError("Unknown model: {}".format(label))
47                    selected_models.add(model)
48                else:
49                    # This is just an app - no model qualifier.
50                    app_label = label
51                    try:
52                        app = apps.get_app_config(app_label)
53                    except LookupError:
54                        raise CommandError("Unknown app: {}".format(app_label))
55                    selected_models.update(app.get_models())
56        for model in selected_models:
57            if is_registered(model):
58                yield model
59