From ec682d41de748adb1718ea5ab87443dd1264ac9c Mon Sep 17 00:00:00 2001
From: Konstantin Alekseev <mail@kalekseev.com>
Date: Mon, 22 Mar 2021 00:32:28 +0300
Subject: [PATCH] Experimental support for autofix

---
 setup.cfg                                     |  5 +-
 shell.nix                                     |  8 ++
 src/extra_checks/checks/base_checks.py        | 34 +++++++-
 src/extra_checks/checks/model_field_checks.py | 61 +++++++++++----
 src/extra_checks/fixes/__init__.py            |  0
 .../fixes/fix_choices_constraint.py           | 65 ++++++++++++++++
 src/extra_checks/management/__init__.py       |  0
 .../management/commands/__init__.py           |  0
 .../management/commands/extra_check.py        | 70 +++++++++++++++++
 tests/test_autofix.py                         | 78 +++++++++++++++++++
 tox.ini                                       |  7 +-
 11 files changed, 304 insertions(+), 24 deletions(-)
 create mode 100644 src/extra_checks/fixes/__init__.py
 create mode 100644 src/extra_checks/fixes/fix_choices_constraint.py
 create mode 100644 src/extra_checks/management/__init__.py
 create mode 100644 src/extra_checks/management/commands/__init__.py
 create mode 100644 src/extra_checks/management/commands/extra_check.py
 create mode 100644 tests/test_autofix.py

diff --git a/setup.cfg b/setup.cfg
index e13b169..8648c89 100644
--- a/setup.cfg
+++ b/setup.cfg
@@ -52,6 +52,9 @@ dev =
     pdbpp
     tox>=3,<4
     black==20.8b1
+    libcst>=0.3,<0.4
+codemod =
+    libcst>=0.3,<0.4
 
 [flake8]
 max-line-length = 110
@@ -64,4 +67,4 @@ force_grid_wrap = 0
 use_parentheses = True
 line_length = 88
 known_first_party = extra_checks
-known_third_party = django,pytest,rest_framework,setuptools
+known_third_party = django,libcst,pytest,rest_framework,setuptools
diff --git a/shell.nix b/shell.nix
index b4c1063..a51caeb 100644
--- a/shell.nix
+++ b/shell.nix
@@ -49,5 +49,13 @@ devshell.mkShell {
       name = "app.lint";
       command = "pre-commit run -a";
     }
+    {
+      help = "run main tox env";
+      name = "app.tox";
+      command = ''
+        unset PYTHONPATH;
+        tox -e 'py{38}-django{22,30,31,32,32-drf,-latest},flake8,black,isort,manifest,mypy,check'
+      '';
+    }
   ];
 }
diff --git a/src/extra_checks/checks/base_checks.py b/src/extra_checks/checks/base_checks.py
index 213b0a0..d1da73e 100644
--- a/src/extra_checks/checks/base_checks.py
+++ b/src/extra_checks/checks/base_checks.py
@@ -15,6 +15,23 @@
 }
 
 
+class ExtraCheckMessage(django.core.checks.CheckMessage):
+    def __init__(
+        self,
+        level: int,
+        msg: str,
+        *,
+        id: str,
+        hint: Optional[str] = None,
+        obj: Any = None,
+        file: Optional[str] = None,
+        fix: Any = None,
+    ) -> None:
+        super().__init__(level, msg, hint=hint, obj=obj, id=id)
+        self._file = file
+        self._fix = fix
+
+
 class BaseCheck(ABC):
     Id: CheckId
     settings_form_class: ClassVar[Type[forms.BaseCheckForm]] = forms.BaseCheckForm
@@ -46,10 +63,21 @@ def is_ignored(self, obj: Any) -> bool:
         return obj in self.ignore_objects or type(obj) in self.ignore_types
 
     def message(
-        self, message: str, hint: Optional[str] = None, obj: Any = None
+        self,
+        message: str,
+        hint: Optional[str] = None,
+        obj: Any = None,
+        file: Optional[str] = None,
+        fix: Any = None,
     ) -> django.core.checks.CheckMessage:
-        return MESSAGE_MAP[self.level](
-            message + f" [{self.Id.value}]", hint=hint, obj=obj, id=self.Id.name
+        return ExtraCheckMessage(
+            self.level,
+            message + f" [{self.Id.value}]",
+            hint=hint,
+            obj=obj,
+            id=self.Id.name,
+            file=file,
+            fix=fix,
         )
 
 
diff --git a/src/extra_checks/checks/model_field_checks.py b/src/extra_checks/checks/model_field_checks.py
index 640b5fa..2397a90 100644
--- a/src/extra_checks/checks/model_field_checks.py
+++ b/src/extra_checks/checks/model_field_checks.py
@@ -285,23 +285,50 @@ def apply(
         field: models.fields.Field,
         ast: FieldASTProtocol,
         model: Type[models.Model],
+        **kwargs: Any,
     ) -> Iterator[django.core.checks.CheckMessage]:
         choices = field.flatchoices  # type: ignore
-        if choices:
-            field_choices = [c[0] for c in choices]
-            if field.blank and "" not in field_choices:
-                field_choices.append("")
-            in_name = f"{field.name}__in"
-            for constraint in model._meta.constraints:
-                if isinstance(constraint, models.CheckConstraint):
-                    conditions = dict(constraint.check.children)
-                    if in_name in conditions and set(field_choices) == set(
-                        conditions[in_name]
-                    ):
-                        return
-            check = f'models.Q({in_name}=[{", ".join([self._repr_choice(c) for c in field_choices])}])'
-            yield self.message(
-                "Field with choices must have companion CheckConstraint to enforce choices on database level.",
-                hint=f'Add to Meta.constraints: `models.CheckConstraint(name="%(app_label)s_%(class)s_{field.name}_valid", check={check})`',
-                obj=field,
+        if not choices:
+            return
+        field_choices = [c[0] for c in choices]
+        if field.blank and "" not in field_choices:
+            field_choices.append("")
+        in_name = f"{field.name}__in"
+        name = f"%(app_label)s_%(class)s_{field.name}_choices_valid"
+        replace = False
+        for constraint in model._meta.constraints:
+            if isinstance(constraint, models.CheckConstraint):
+                if name == constraint.name:
+                    replace = True
+                conditions = dict(constraint.check.children)
+                if in_name in conditions and set(field_choices) == set(
+                    conditions[in_name]
+                ):
+                    return
+        check = f'models.Q({in_name}=[{", ".join([self._repr_choice(c) for c in field_choices])}])'
+        kwargs = {}
+        try:
+            import importlib
+
+            from extra_checks.fixes.fix_choices_constraint import (
+                gen_fix_for_choices_constraint,
             )
+        except ImportError:
+            pass
+        else:
+            kwargs = {
+                "file": importlib.import_module(model.__module__).__file__,
+                "fix": gen_fix_for_choices_constraint(
+                    model.__name__,
+                    f"%(app_label)s_%(class)s_{field.name}_choices_valid",
+                    check=check,
+                    replace=replace,
+                ),
+            }
+
+        yield self.message(
+            "Field with choices must have companion CheckConstraint to enforce choices on database level.",
+            hint=f'Add to Meta.constraints: `models.CheckConstraint(name="{name}", check={check})`',
+            obj=field,
+            **kwargs,
+        )
diff --git a/src/extra_checks/fixes/__init__.py b/src/extra_checks/fixes/__init__.py
new file mode 100644
index 0000000..e69de29
diff --git a/src/extra_checks/fixes/fix_choices_constraint.py b/src/extra_checks/fixes/fix_choices_constraint.py
new file mode 100644
index 0000000..7249e4b
--- /dev/null
+++ b/src/extra_checks/fixes/fix_choices_constraint.py
@@ -0,0 +1,65 @@
+import libcst as cst
+from libcst import matchers as m
+
+
+def gen_fix_for_choices_constraint(
+    class_name: str, name: str, check: str, replace: bool = False
+) -> m.MatcherDecoratableTransformer:
+    class Fixes(m.MatcherDecoratableTransformer):
+        def __init__(self) -> None:
+            self.is_constraint_updated = False
+            super().__init__()
+
+        @m.call_if_inside(m.ClassDef(m.Name(class_name)))
+        @m.leave(m.ClassDef(m.Name("Meta")))
+        def leave_meta(
+            self, node: cst.ClassDef, updated_node: cst.ClassDef
+        ) -> cst.ClassDef:
+            if not self.is_constraint_updated and not replace:
+                exp = cst.parse_statement(
+                    f'constraints = [models.CheckConstraint(name="{name}", check={check})]'
+                )
+                lines = updated_node.body.body
+                return updated_node.with_deep_changes(
+                    updated_node.body, body=[*lines, exp]
+                )
+            self.is_constraint_updated = False
+            return updated_node
+
+        if replace:
+
+            @m.call_if_inside(m.ClassDef(m.Name(class_name)))
+            @m.call_if_inside(m.ClassDef(m.Name("Meta")))
+            @m.call_if_inside(
+                m.Assign(targets=[m.AssignTarget(target=m.Name("constraints"))])
+            )
+            @m.leave(
+                m.Call(
+                    func=m.Attribute(attr=m.Name("CheckConstraint"))
+                    | m.Name("CheckConstraint")
+                )
+            )
+            def fix_existing(self, node: cst.Call, updated_node: cst.Call) -> cst.Call:
+                # TODO select either models.CheckConstraint or CheckConstraint
+                if node.args[0].value.raw_value == name:
+                    return cst.parse_expression(
+                        f'models.CheckConstraint(name="{name}", check={check})'
+                    )
+                return updated_node
+
+        else:
+
+            @m.call_if_inside(m.ClassDef(m.Name(class_name)))
+            @m.call_if_inside(m.ClassDef(m.Name("Meta")))
+            @m.leave(m.Assign(targets=[m.AssignTarget(target=m.Name("constraints"))]))
+            def add_new(self, node: cst.Assign, updated_node: cst.Assign) -> cst.Assign:
+                self.is_constraint_updated = True
+                exp = cst.parse_expression(
+                    f'[models.CheckConstraint(name="{name}", check={check})]'
+                )
+                constraints = updated_node.value.elements
+                return updated_node.with_deep_changes(
+                    updated_node.value, elements=[*constraints, exp.elements[0]]
+                )
+
+    return Fixes()
diff --git a/src/extra_checks/management/__init__.py b/src/extra_checks/management/__init__.py
new file mode 100644
index 0000000..e69de29
diff --git a/src/extra_checks/management/commands/__init__.py b/src/extra_checks/management/commands/__init__.py
new file mode 100644
index 0000000..e69de29
diff --git a/src/extra_checks/management/commands/extra_check.py b/src/extra_checks/management/commands/extra_check.py
new file mode 100644
index 0000000..bd6b399
--- /dev/null
+++ b/src/extra_checks/management/commands/extra_check.py
@@ -0,0 +1,70 @@
+from typing import TYPE_CHECKING, Any, Dict, List
+
+from django.core import checks
+from django.core.management.base import SystemCheckError
+from django.core.management.commands.check import Command as BaseCommand
+
+if TYPE_CHECKING:
+    from extra_checks.checks.base_checks import ExtraCheckMessage
+
+
+class Command(BaseCommand):
+    def __init__(self, *args: Any, **kwargs: Any) -> None:
+        super().__init__(*args, **kwargs)
+        self._errors: "List[ExtraCheckMessage]" = []
+
+    def add_arguments(self, parser: Any) -> None:
+        super().add_arguments(parser)
+        parser.add_argument(
+            "--fix",
+            action="store_true",
+            help="Apply autofix if available.",
+        )
+        parser.add_argument(
+            "--fix-black",
+            action="store_true",
+            help="Apply black on autofix result.",
+        )
+
+    def _run_checks(self, **kwargs: Any) -> dict:
+        errors = checks.run_checks(**kwargs)
+        self._errors = errors
+        return errors
+
+    def handle(self, *app_labels: Any, **options: Any) -> None:
+        on_exit = None
+        try:
+            super().handle(*app_labels, **options)
+        except SystemCheckError as exc:
+            on_exit = exc
+        if not options["fix"]:
+            return
+        import libcst as cst
+        from libcst import matchers as m
+
+        files: Dict[str, List[m.MatcherDecoratableTransformer]] = {}
+        for error in self._errors:
+            fix = getattr(error, "_fix", None)
+            file = getattr(error, "_file", None)
+            if fix and file:
+                files.setdefault(file, []).append(fix)
+        for file, fixes in files.items():
+            with open(file, "r") as f:
+                source_text = f.read()
+            tree = cst.parse_module(source_text)
+            for fix in fixes:
+                tree = tree.visit(fix)
+            result_text = tree.code
+            if source_text != result_text:
+                if options["fix_black"]:
+                    import black
+
+                    mode = black.FileMode()
+                    fast = False
+                    result_text = black.format_file_contents(
+                        src_contents=result_text, fast=fast, mode=mode
+                    )
+                with open(file, "w") as f:
+                    f.write(result_text)
+        if on_exit:
+            raise on_exit
diff --git a/tests/test_autofix.py b/tests/test_autofix.py
new file mode 100644
index 0000000..d75d5ff
--- /dev/null
+++ b/tests/test_autofix.py
@@ -0,0 +1,78 @@
+import libcst as cst
+
+from extra_checks.fixes.fix_choices_constraint import gen_fix_for_choices_constraint
+
+SOURCE = """
+class TestClass(Token):
+    class Meta:
+        constraints = [
+            models.CheckConstraint(name="value_valid", check=models.Q(value__in=[1, 2]))
+        ]
+    """
+
+
+def test_fix_add():
+    result = """
+class TestClass(Token):
+    class Meta:
+        constraints = [
+            models.CheckConstraint(name="value_valid", check=models.Q(value__in=[1, 2])), models.CheckConstraint(name="another_valid", check=models.Q(value__in=[1, 2, 3]))
+        ]
+    """
+
+    source_tree = cst.parse_module(SOURCE)
+    modefied_tree = source_tree.visit(
+        gen_fix_for_choices_constraint(
+            "TestClass",
+            name="another_valid",
+            check="models.Q(value__in=[1, 2, 3])",
+            replace=False,
+        )
+    )
+    assert modefied_tree.code == result
+
+
+def test_fix_replace():
+    result = """
+class TestClass(Token):
+    class Meta:
+        constraints = [
+            models.CheckConstraint(name="value_valid", check=models.Q(value__in=[1, 2, 3]))
+        ]
+    """
+
+    source_tree = cst.parse_module(SOURCE)
+    modefied_tree = source_tree.visit(
+        gen_fix_for_choices_constraint(
+            "TestClass",
+            name="value_valid",
+            check="models.Q(value__in=[1, 2, 3])",
+            replace=True,
+        )
+    )
+    assert modefied_tree.code == result
+
+
+def test_fix_meta_add_constraints():
+    source = """
+class TestClass(Token):
+    class Meta:
+        db_table = 'test_class'
+    """
+    result = """
+class TestClass(Token):
+    class Meta:
+        db_table = 'test_class'
+        constraints = [models.CheckConstraint(name="value_valid", check=models.Q(value__in=[1, 2, 3]))]
+    """
+
+    source_tree = cst.parse_module(source)
+    modefied_tree = source_tree.visit(
+        gen_fix_for_choices_constraint(
+            "TestClass",
+            name="value_valid",
+            check="models.Q(value__in=[1, 2, 3])",
+            replace=False,
+        )
+    )
+    assert modefied_tree.code == result
diff --git a/tox.ini b/tox.ini
index 8ccb4a6..ff1e8e6 100644
--- a/tox.ini
+++ b/tox.ini
@@ -17,9 +17,10 @@ setenv =
 usedevelop = false
 deps =
     py36,py37: typing_extensions
-    pytest
-    pytest-django
-    pytest-cov
+    pytest>=6,<7
+    pytest-cov>=2,<3
+    pytest-django>=4,<5
+    libcst>=0.3,<0.4
     django22: Django>=2.2,<2.3
     django30: Django>=3.0,<3.1
     django31: Django>=3.1,<3.2