Add internal classes and add type hints + mypy.

django_migration_linter/sql_analyser/__init__.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 from .base import BaseAnalyser  # noqa
 from .mysql import MySqlAnalyser  # noqa
 from .postgresql import PostgresqlAnalyser  # noqa

django_migration_linter/sql_analyser/analyser.py

@@ -1,6 +1,13 @@
+from __future__ import annotations
+
 import logging
+from typing import TYPE_CHECKING, Iterable, Type
+
+if TYPE_CHECKING:
+    from sql_analyser.base import Issue
 
 from django_migration_linter.sql_analyser import (
+    BaseAnalyser,
     MySqlAnalyser,
     PostgresqlAnalyser,
     SqliteAnalyser,
@@ -8,20 +15,22 @@ from django_migration_linter.sql_analyser import (
 
 logger = logging.getLogger("django_migration_linter")
 
-ANALYSER_STRING_MAPPING = {
+ANALYSER_STRING_MAPPING: dict[str, Type[BaseAnalyser]] = {
     "sqlite": SqliteAnalyser,
     "mysql": MySqlAnalyser,
     "postgresql": PostgresqlAnalyser,
 }
 
 
-def get_sql_analyser_class(database_vendor, analyser_string=None):
+def get_sql_analyser_class(
+    database_vendor: str, analyser_string: str | None = None
+) -> Type[BaseAnalyser]:
     if analyser_string:
         return get_sql_analyser_from_string(analyser_string)
     return get_sql_analyser_class_from_db_vendor(database_vendor)
 
 
-def get_sql_analyser_from_string(analyser_string):
+def get_sql_analyser_from_string(analyser_string: str) -> Type[BaseAnalyser]:
     if analyser_string not in ANALYSER_STRING_MAPPING:
         raise ValueError(
             "Unknown SQL analyser '{}'. Known values: '{}'".format(
@@ -32,7 +41,8 @@ def get_sql_analyser_from_string(analyser_string):
     return ANALYSER_STRING_MAPPING[analyser_string]
 
 
-def get_sql_analyser_class_from_db_vendor(database_vendor):
+def get_sql_analyser_class_from_db_vendor(database_vendor: str) -> Type[BaseAnalyser]:
+    sql_analyser_class: Type[BaseAnalyser]
     if "mysql" in database_vendor:
         sql_analyser_class = MySqlAnalyser
     elif "postgre" in database_vendor:
@@ -51,8 +61,10 @@ def get_sql_analyser_class_from_db_vendor(database_vendor):
 
 
 def analyse_sql_statements(
-    sql_analyser_class, sql_statements, exclude_migration_tests=None
-):
+    sql_analyser_class: Type[BaseAnalyser],
+    sql_statements: list[str],
+    exclude_migration_tests: Iterable[str] | None = None,
+) -> tuple[list[Issue], list[Issue], list[Issue]]:
     sql_analyser = sql_analyser_class(exclude_migration_tests)
     sql_analyser.analyse(sql_statements)
     return sql_analyser.errors, sql_analyser.ignored, sql_analyser.warnings

django_migration_linter/sql_analyser/base.py

@@ -1,12 +1,36 @@
+from __future__ import annotations
+
 import logging
 import re
-
-from .utils import update_migration_tests
+from copy import deepcopy
+from dataclasses import dataclass
+from enum import Enum
+from typing import Callable, Iterable
 
 logger = logging.getLogger("django_migration_linter")
 
 
-def has_not_null_column(sql_statements, **kwargs):
+def find_check_from_code(checks: list[Check], code: str) -> Check | None:
+    return next((c for c in checks if c.code == code), None)
+
+
+def update_migration_checks(
+    base_checks: list[Check], specific_checks: list[Check]
+) -> list[Check]:
+    new_checks = deepcopy(base_checks)
+    for override_check in specific_checks:
+        migration_check = find_check_from_code(new_checks, override_check.code)
+
+        if migration_check is None or not override_check.code:
+            migration_check = deepcopy(override_check)
+            new_checks.append(migration_check)
+
+        for field in override_check.__dataclass_fields__:
+            setattr(migration_check, field, getattr(override_check, field))
+    return new_checks
+
+
+def has_not_null_column(sql_statements: list[str], **kwargs) -> bool:
     # TODO: improve to detect that the same column is concerned
     ends_with_default = False
     for sql in sql_statements:
@@ -23,7 +47,7 @@ def has_not_null_column(sql_statements, **kwargs):
     )
 
 
-def has_add_unique(sql_statements, **kwargs):
+def has_add_unique(sql_statements: list[str], **kwargs) -> bool:
     regex_result = None
     for sql in sql_statements:
         regex_result = re.search(
@@ -41,130 +65,167 @@ def has_add_unique(sql_statements, **kwargs):
     return not table_is_added_in_transaction
 
 
+class CheckMode(Enum):
+    """
+    Defines whether the Check.fn gets a str or a list[str] as first parameter.
+    """
+
+    ONE_LINER = 1
+    TRANSACTION = 2
+
+
+class CheckType(Enum):
+    ERROR = 1
+    WARNING = 2
+
+
+@dataclass
+class Check:
+    """
+    Represents a check that will be done against SQL statement(s).
+    """
+
+    code: str
+    fn: Callable
+    message: str
+    mode: CheckMode
+    type: CheckType
+
+
+@dataclass
+class Issue:
+    code: str
+    message: str
+    table: str | None = None
+    column: str | None = None
+
+
 class BaseAnalyser:
-    base_migration_tests = [
-        {
-            "code": "RENAME_TABLE",
-            "fn": lambda sql, **kw: re.search("RENAME TABLE", sql)
+    base_migration_checks: list[Check] = [
+        Check(
+            code="RENAME_TABLE",
+            fn=lambda sql, **kw: re.search("RENAME TABLE", sql)
             or re.search("ALTER TABLE .* RENAME TO", sql),
-            "msg": "RENAMING tables",
-            "mode": "one_liner",
-            "type": "error",
-        },
-        {
-            "code": "NOT_NULL",
-            "fn": has_not_null_column,
-            "msg": "NOT NULL constraint on columns",
-            "mode": "transaction",
-            "type": "error",
-        },
-        {
-            "code": "DROP_COLUMN",
-            "fn": lambda sql, **kw: re.search("DROP COLUMN", sql),
-            "msg": "DROPPING columns",
-            "mode": "one_liner",
-            "type": "error",
-        },
-        {
-            "code": "DROP_TABLE",
-            "fn": lambda sql, **kw: sql.startswith("DROP TABLE"),
-            "msg": "DROPPING table",
-            "mode": "one_liner",
-            "type": "error",
-        },
-        {
-            "code": "RENAME_COLUMN",
-            "fn": lambda sql, **kw: re.search("ALTER TABLE .* CHANGE", sql)
+            message="RENAMING tables",
+            mode=CheckMode.ONE_LINER,
+            type=CheckType.ERROR,
+        ),
+        Check(
+            code="NOT_NULL",
+            fn=has_not_null_column,
+            message="NOT NULL constraint on columns",
+            mode=CheckMode.TRANSACTION,
+            type=CheckType.ERROR,
+        ),
+        Check(
+            code="DROP_COLUMN",
+            fn=lambda sql, **kw: re.search("DROP COLUMN", sql),
+            message="DROPPING columns",
+            mode=CheckMode.ONE_LINER,
+            type=CheckType.ERROR,
+        ),
+        Check(
+            code="DROP_TABLE",
+            fn=lambda sql, **kw: sql.startswith("DROP TABLE"),
+            message="DROPPING table",
+            mode=CheckMode.ONE_LINER,
+            type=CheckType.ERROR,
+        ),
+        Check(
+            code="RENAME_COLUMN",
+            fn=lambda sql, **kw: re.search("ALTER TABLE .* CHANGE", sql)
             or re.search("ALTER TABLE .* RENAME COLUMN", sql),
-            "msg": "RENAMING columns",
-            "mode": "one_liner",
-            "type": "error",
-        },
-        {
-            "code": "ALTER_COLUMN",
-            "fn": lambda sql, **kw: re.search(
-                "ALTER TABLE .* ALTER COLUMN .* TYPE", sql
-            ),
-            "msg": (
+            message="RENAMING columns",
+            mode=CheckMode.ONE_LINER,
+            type=CheckType.ERROR,
+        ),
+        Check(
+            code="ALTER_COLUMN",
+            fn=lambda sql, **kw: re.search("ALTER TABLE .* ALTER COLUMN .* TYPE", sql),
+            message=(
                 "ALTERING columns (Could be backward compatible. "
                 "You may ignore this migration.)"
             ),
-            "mode": "one_liner",
-            "type": "error",
-        },
-        {
-            "code": "ADD_UNIQUE",
-            "fn": has_add_unique,
-            "msg": "ADDING unique constraint",
-            "mode": "transaction",
-            "type": "error",
-        },
+            mode=CheckMode.ONE_LINER,
+            type=CheckType.ERROR,
+        ),
+        Check(
+            code="ADD_UNIQUE",
+            fn=has_add_unique,
+            message="ADDING unique constraint",
+            mode=CheckMode.TRANSACTION,
+            type=CheckType.ERROR,
+        ),
     ]
 
-    migration_tests = []
+    migration_checks: list[Check] = []
 
-    def __init__(self, exclude_migration_tests):
-        self.exclude_migration_tests = exclude_migration_tests or []
-        self.errors = []
-        self.warnings = []
-        self.ignored = []
-        self.migration_tests = update_migration_tests(
-            self.base_migration_tests, self.migration_tests
+    def __init__(self, exclude_migration_tests: Iterable[str] | None):
+        self.exclude_migration_tests: Iterable[str] = exclude_migration_tests or []
+        self.errors: list[Issue] = []
+        self.warnings: list[Issue] = []
+        self.ignored: list[Issue] = []
+        self.migration_checks = update_migration_checks(
+            self.base_migration_checks, self.migration_checks
         )
 
-    def analyse(self, sql_statements):
+    def analyse(self, sql_statements: list[str]) -> None:
         for statement in sql_statements:
-            for test in self.one_line_migration_tests:
-                self._test_sql(test, sql=statement)
+            for test in self.one_line_migration_checks:
+                self._check_sql(test, sql=statement)
 
-        for test in self.transaction_migration_tests:
-            self._test_sql(test, sql=sql_statements)
+        for test in self.transaction_migration_checks:
+            self._check_sql(test, sql=sql_statements)
 
     @property
-    def one_line_migration_tests(self):
-        return (test for test in self.migration_tests if test["mode"] == "one_liner")
+    def one_line_migration_checks(self) -> Iterable[Check]:
+        return (c for c in self.migration_checks if c.mode == CheckMode.ONE_LINER)
 
     @property
-    def transaction_migration_tests(self):
-        return (test for test in self.migration_tests if test["mode"] == "transaction")
+    def transaction_migration_checks(self) -> Iterable[Check]:
+        return (c for c in self.migration_checks if c.mode == CheckMode.TRANSACTION)
 
-    def _test_sql(self, test, sql):
-        if test["fn"](sql, errors=self.errors):
-            if test["code"] in self.exclude_migration_tests:
+    def _check_sql(self, check: Check, sql: list[str] | str) -> None:
+        if check.fn(sql, errors=self.errors):
+            if check.code in self.exclude_migration_tests:
                 action = "IGNORED"
                 list_to_add = self.ignored
-            elif test["type"] == "warning":
+            elif check.type == CheckType.WARNING:
                 action = "WARNING"
                 list_to_add = self.warnings
             else:
                 action = "ERROR"
                 list_to_add = self.errors
             logger.debug("Testing %s -- %s", sql, action)
-            err = self.build_error_dict(migration_test=test, sql_statement=sql)
-            list_to_add.append(err)
+            issue = self.build_issue(migration_check=check, sql_statement=sql)
+            list_to_add.append(issue)
         else:
             logger.debug("Testing %s -- PASSED", sql)
 
-    def build_error_dict(self, migration_test, sql_statement):
+    def build_issue(
+        self, migration_check: Check, sql_statement: list[str] | str
+    ) -> Issue:
         table = self.detect_table(sql_statement)
         col = self.detect_column(sql_statement)
-        return {
-            "msg": migration_test["msg"],
-            "code": migration_test["code"],
-            "table": table,
-            "column": col,
-        }
+        return Issue(
+            code=migration_check.code,
+            message=migration_check.message,
+            table=table,
+            column=col,
+        )
 
     @staticmethod
-    def detect_table(sql):
+    def detect_table(sql: list[str] | str) -> str | None:
         if isinstance(sql, str):
             regex_result = re.search("TABLE [`\"'](.*?)[`\"']", sql, re.IGNORECASE)
             if regex_result:
                 return regex_result.group(1)
+        return None
 
     @staticmethod
-    def detect_column(sql):
+    def detect_column(sql: list[str] | str) -> str | None:
         if isinstance(sql, str):
             regex_result = re.search("COLUMN [`\"'](.*?)[`\"']", sql, re.IGNORECASE)
             if regex_result:
                 return regex_result.group(1)
+        return None

django_migration_linter/sql_analyser/mysql.py

@@ -1,22 +1,26 @@
+from __future__ import annotations
+
 import re
 
-from .base import BaseAnalyser
+from .base import BaseAnalyser, Check, CheckMode, CheckType
 
 
 class MySqlAnalyser(BaseAnalyser):
-    migration_tests = [
-        {
-            "code": "ALTER_COLUMN",
-            "fn": lambda sql, **kw: re.search(
-                "ALTER TABLE .* MODIFY .* (?!NULL);?$", sql
+    migration_checks: list[Check] = [
+        Check(
+            code="ALTER_COLUMN",
+            fn=lambda sql, **kw: re.search("ALTER TABLE .* MODIFY .* (?!NULL);?$", sql),
+            message=(
+                "ALTERING columns (Could be backward compatible. "
+                "You may ignore this migration.)"
             ),
-            "mode": "one_liner",
-            "type": "error",
-        }
+            mode=CheckMode.ONE_LINER,
+            type=CheckType.ERROR,
+        ),
     ]
 
     @staticmethod
-    def detect_column(sql):
+    def detect_column(sql: list[str] | str) -> str | None:
         if isinstance(sql, str):
             regex_result = re.search("COLUMN [`\"'](.*?)[`\"']", sql, re.IGNORECASE)
             if regex_result:
@@ -24,3 +28,4 @@ class MySqlAnalyser(BaseAnalyser):
             regex_result = re.search("MODIFY [`\"'](.*?)[`\"']", sql, re.IGNORECASE)
             if regex_result:
                 return regex_result.group(1)
+        return None

django_migration_linter/sql_analyser/postgresql.py

@@ -1,9 +1,11 @@
+from __future__ import annotations
+
 import re
 
-from .base import BaseAnalyser
+from .base import BaseAnalyser, Check, CheckMode, CheckType
 
 
-def has_create_index(sql_statements, **kwargs):
+def has_create_index(sql_statements: list[str], **kwargs) -> bool:
     regex_result = None
     for sql in sql_statements:
         regex_result = re.search(r"CREATE (UNIQUE )?INDEX.*ON (.*) \(", sql)
@@ -22,27 +24,27 @@ def has_create_index(sql_statements, **kwargs):
 
 
 class PostgresqlAnalyser(BaseAnalyser):
-    migration_tests = [
-        {
-            "code": "CREATE_INDEX",
-            "fn": has_create_index,
-            "msg": "CREATE INDEX locks table",
-            "mode": "transaction",
-            "type": "warning",
-        },
-        {
-            "code": "DROP_INDEX",
-            "fn": lambda sql, **kw: re.search("DROP INDEX", sql)
+    migration_checks: list[Check] = [
+        Check(
+            code="CREATE_INDEX",
+            fn=has_create_index,
+            message="CREATE INDEX locks table",
+            mode=CheckMode.TRANSACTION,
+            type=CheckType.WARNING,
+        ),
+        Check(
+            code="DROP_INDEX",
+            fn=lambda sql, **kw: re.search("DROP INDEX", sql)
             and not re.search("INDEX CONCURRENTLY", sql),
-            "msg": "DROP INDEX locks table",
-            "mode": "one_liner",
-            "type": "warning",
-        },
-        {
-            "code": "REINDEX",
-            "fn": lambda sql, **kw: sql.startswith("REINDEX"),
-            "msg": "REINDEX locks table",
-            "mode": "one_liner",
-            "type": "warning",
-        },
+            message="DROP INDEX locks table",
+            mode=CheckMode.ONE_LINER,
+            type=CheckType.WARNING,
+        ),
+        Check(
+            code="REINDEX",
+            fn=lambda sql, **kw: sql.startswith("REINDEX"),
+            message="REINDEX locks table",
+            mode=CheckMode.ONE_LINER,
+            type=CheckType.WARNING,
+        ),
     ]

django_migration_linter/sql_analyser/sqlite.py

@@ -1,31 +1,35 @@
+from __future__ import annotations
+
 import re
 
-from .base import BaseAnalyser
+from .base import BaseAnalyser, Check, CheckMode, CheckType
 
 
 class SqliteAnalyser(BaseAnalyser):
-    migration_tests = [
-        {
-            "code": "RENAME_TABLE",
-            "fn": lambda sql, **kw: re.search("ALTER TABLE .* RENAME TO", sql)
+    migration_checks: list[Check] = [
+        Check(
+            code="RENAME_TABLE",
+            fn=lambda sql, **kw: re.search("ALTER TABLE .* RENAME TO", sql)
             and "__old" not in sql
             and "new__" not in sql,
-            "type": "error",
-        },
-        {
-            "code": "DROP_TABLE",
+            message="RENAMING tables",
+            mode=CheckMode.ONE_LINER,
+            type=CheckType.ERROR,
+        ),
+        Check(
+            code="DROP_TABLE",
             # TODO: improve to detect that the table names overlap
-            "fn": lambda sql_statements, **kw: any(
+            fn=lambda sql_statements, **kw: any(
                 sql.startswith("DROP TABLE") for sql in sql_statements
             )
             and not any(sql.startswith("CREATE TABLE") for sql in sql_statements),
-            "msg": "DROPPING table",
-            "mode": "transaction",
-            "type": "error",
-        },
-        {
-            "code": "NOT_NULL",
-            "fn": lambda sql_statements, **kw: any(
+            message="DROPPING table",
+            mode=CheckMode.TRANSACTION,
+            type=CheckType.ERROR,
+        ),
+        Check(
+            code="NOT_NULL",
+            fn=lambda sql_statements, **kw: any(
                 re.search("NOT NULL(?! PRIMARY)(?! DEFAULT)", sql)
                 for sql in sql_statements
             )
@@ -34,13 +38,14 @@ class SqliteAnalyser(BaseAnalyser):
                 and ("__old" in sql or "new__" in sql)
                 for sql in sql_statements
             ),
-            "mode": "transaction",
-            "type": "error",
-        },
+            message="NOT NULL constraint on columns",
+            mode=CheckMode.TRANSACTION,
+            type=CheckType.ERROR,
+        ),
     ]
 
     @staticmethod
-    def detect_table(sql):
+    def detect_table(sql: list[str] | str) -> str | None:
         if isinstance(sql, str):
             regex_result = re.search("TABLE [`\"'](.*?)[`\"']", sql, re.IGNORECASE)
             if regex_result:
@@ -48,3 +53,4 @@ class SqliteAnalyser(BaseAnalyser):
             regex_result = re.search("ON [`\"'](.*?)[`\"']", sql, re.IGNORECASE)
             if regex_result:
                 return regex_result.group(1)
+        return None

django_migration_linter/sql_analyser/utils.py

@@ -1,21 +0,0 @@
-from copy import deepcopy
-
-
-def find_error_dict_with_code(tests, code):
-    return next((test_dict for test_dict in tests if test_dict["code"] == code), None)
-
-
-def update_migration_tests(base_tests, specific_tests):
-    base_tests = deepcopy(base_tests)
-    for override_test in specific_tests:
-        migration_test_dict = find_error_dict_with_code(
-            base_tests, override_test["code"]
-        )
-
-        if migration_test_dict is None or not override_test["code"]:
-            migration_test_dict = {}
-            base_tests.append(migration_test_dict)
-
-        for key in override_test.keys():
-            migration_test_dict[key] = override_test[key]
-    return base_tests