Adapt 'warnings-as-errors' option to allow taking migration test codes as input

django_migration_linter/management/commands/lintmigrations.py

@@ -11,7 +11,11 @@ from django.core.management.base import BaseCommand
 from ...constants import __version__
 from ...migration_linter import MessageType, MigrationLinter
 from ...sql_analyser.analyser import ANALYSER_STRING_MAPPING
-from ..utils import configure_logging, register_linting_configuration_options
+from ..utils import (
+    configure_logging,
+    extract_warnings_as_errors_option,
+    register_linting_configuration_options,
+)
 
 CONFIG_NAME = "django_migration_linter"
 PYPROJECT_TOML = "pyproject.toml"
@@ -155,6 +159,11 @@ class Command(BaseCommand):
             if not options[k]:
                 options[k] = v
 
+        (
+            warnings_as_errors_tests,
+            all_warnings_as_errors,
+        ) = extract_warnings_as_errors_option(options["warnings_as_errors"])
+
         linter = MigrationLinter(
             settings_path,
             ignore_name_contains=options["ignore_name_contains"],
@@ -170,7 +179,8 @@ class Command(BaseCommand):
             only_unapplied_migrations=options["unapplied_migrations"],
             exclude_migration_tests=options["exclude_migration_tests"],
             quiet=options["quiet"],
-            warnings_as_errors=options["warnings_as_errors"],
+            warnings_as_errors_tests=warnings_as_errors_tests,
+            all_warnings_as_errors=all_warnings_as_errors,
             no_output=options["verbosity"] == 0,
             analyser_string=options["sql_analyser"],
         )

django_migration_linter/management/commands/makemigrations.py

@@ -9,7 +9,11 @@ from django.db.migrations.writer import MigrationWriter
 
 from django_migration_linter import MigrationLinter
 
-from ..utils import configure_logging, register_linting_configuration_options
+from ..utils import (
+    configure_logging,
+    extract_warnings_as_errors_option,
+    register_linting_configuration_options,
+)
 
 
 def ask_should_keep_migration():
@@ -68,13 +72,19 @@ class Command(MakeMigrationsCommand):
             else default_should_keep_migration
         )
 
+        (
+            warnings_as_errors_tests,
+            all_warnings_as_errors,
+        ) = extract_warnings_as_errors_option(self.warnings_as_errors)
+
         # Lint migrations
         linter = MigrationLinter(
             path=os.environ["DJANGO_SETTINGS_MODULE"],
             database=self.database,
             no_cache=True,
             exclude_migration_tests=self.exclude_migrations_tests,
-            warnings_as_errors=self.warnings_as_errors,
+            warnings_as_errors_tests=warnings_as_errors_tests,
+            all_warnings_as_errors=all_warnings_as_errors,
         )
 
         for app_label, app_migrations in changes.items():

django_migration_linter/management/utils.py

@@ -21,8 +21,10 @@ def register_linting_configuration_options(parser):
 
     parser.add_argument(
         "--warnings-as-errors",
-        action="store_true",
-        help="handle warnings as errors",
+        type=str,
+        nargs="*",
+        help="handle warnings as errors. Optionally specify the tests to handle as "
+        "errors (e.g. RUNPYTHON_REVERSIBLE)",
     )
 
 
@@ -35,3 +37,16 @@ def configure_logging(verbosity):
         logger.disabled = True
     else:
         logging.basicConfig(format="%(message)s")
+
+
+def extract_warnings_as_errors_option(warnings_as_errors):
+    if isinstance(warnings_as_errors, list):
+        warnings_as_errors_tests = warnings_as_errors
+        # If the option is specified but without any test codes,
+        # all warnings become errors
+        all_warnings_as_errors = len(warnings_as_errors) == 0
+    else:
+        warnings_as_errors_tests = None
+        all_warnings_as_errors = False
+
+    return warnings_as_errors_tests, all_warnings_as_errors