Adds ability to rename elidable functions found in operations

django_squash/db/migrations/autodetector.py

@@ -84,10 +84,13 @@ def add_non_elidables(self, loader, changes):
                         continue
 
                     if isinstance(operation, dj_migrations.RunSQL):
+                        operation._original_migration = migration
                         new_operations.append(operation)
                     elif isinstance(operation, dj_migrations.RunPython):
+                        operation._original_migration = migration
                         new_operations.append(operation)
                     elif isinstance(operation, postgres.PGCreateExtension):
+                        operation._original_migration = migration
                         new_operations_bubble_top.append(operation)
                     elif isinstance(operation, dj_migrations.SeparateDatabaseAndState):
                         # A valid use case for this should be given before any work is done.

django_squash/db/migrations/utils.py

@@ -1,4 +1,5 @@
 import ast
+import functools
 import hashlib
 import importlib
 import inspect
@@ -12,6 +13,19 @@
 from django.db import migrations
 from django.utils.module_loading import import_string
 
+from django_squash import settings as app_settings
+
+
+@functools.lru_cache(maxsize=1)
+def get_custom_rename_function():
+    """
+    Custom function naming when copying elidable functions from one file to another.
+    """
+    custom_rename_function = app_settings.DJANGO_SQUASH_CUSTOM_RENAME_FUNCTION
+
+    if custom_rename_function:
+        return import_string(custom_rename_function)
+
 
 def file_hash(file_path):
     """
@@ -38,9 +52,14 @@ class UniqueVariableName:
     This class will return a unique name for a variable / function.
     """
 
-    def __init__(self):
+    def __init__(self, context, naming_function=None):
         self.names = defaultdict(int)
         self.functions = {}
+        self.context = context
+        self.naming_function = naming_function or (lambda n, c: n)
+
+    def update_context(self, context):
+        self.context.update(context)
 
     def function(self, func):
         if not callable(func):
@@ -52,39 +71,37 @@ def function(self, func):
         if inspect.ismethod(func) or inspect.signature(func).parameters.get("self") is not None:
             raise ValueError("func cannot be part of an instance")
 
-        name = original_name = func.__qualname__
+        name = func.__qualname__
         if "." in name:
             parent_name, actual_name = name.rsplit(".", 1)
             parent = getattr(import_string(func.__module__), parent_name)
             if issubclass(parent, migrations.Migration):
-                name = name = original_name = actual_name
-        already_accounted = func in self.functions
-        if already_accounted:
+                name = actual_name
+
+        if func in self.functions:
             return self.functions[func]
 
+        name = self.naming_function(name, {**self.context, "type_": "function", "func": func})
+        new_name = self.functions[func] = self.uniq(name)
+
+        return new_name
+
+    def uniq(self, name, original_name=None):
+        original_name = original_name or name
         # Endless loop that will try different combinations until it finds a unique name
         for i, _ in enumerate(itertools.count(), 2):
             if self.names[name] == 0:
-                self.functions[func] = name
                 self.names[name] += 1
                 break
 
             name = "%s_%s" % (original_name, i)
-
-        self.functions[func] = name
-
         return name
 
     def __call__(self, name, force_number=False):
-        self.names[name] += 1
-        count = self.names[name]
-        if not force_number and count == 1:
-            return name
-        else:
-            new_name = "%s_%s" % (name, count)
-            # Make sure that the function name is fully unique
-            # You can potentially have the same name already defined.
-            return self(new_name)
+        original_name = name
+        if force_number:
+            name = f"{name}_1"
+        return self.uniq(name, original_name)
 
 
 def get_imports(module):
@@ -109,9 +126,10 @@ def get_imports(module):
 
 
 def normalize_function_name(name):
-    class_name, _, function_name = name.rpartition(".")
-    if class_name and not function_name:
-        function_name = class_name
+    _, _, function_name = name.rpartition(".")
+    if function_name[0].isdigit():
+        # Functions CANNOT start with a number
+        function_name = "f_" + function_name
     return function_name
 
 
@@ -123,10 +141,7 @@ def copy_func(f, name):
     func.__qualname__ = name
     func.__original__ = f
     func.__source__ = re.sub(
-        rf"(def\s+){normalize_function_name(f.__qualname__)}",
-        rf"\1{normalize_function_name(name)}",
-        inspect.getsource(f),
-        1,
+        rf"(def\s+){normalize_function_name(f.__qualname__)}", rf"\1{name}", inspect.getsource(f), 1
     )
     return func
 

django_squash/db/migrations/writer.py

@@ -166,8 +166,20 @@ def as_string(self):
             return self.replace_in_migration()
 
         variables = []
-        unique_names = utils.UniqueVariableName()
+        custom_naming_function = utils.get_custom_rename_function()
+        unique_names = utils.UniqueVariableName(
+            {"app": self.migration.app_label}, naming_function=custom_naming_function
+        )
         for operation in self.migration.operations:
+            unique_names.update_context(
+                {
+                    "new_migration": self.migration,
+                    "operation": operation,
+                    "migration": (
+                        operation._original_migration if hasattr(operation, "_original_migration") else self.migration
+                    ),
+                }
+            )
             operation._deconstruct = operation.__class__.deconstruct
 
             def deconstruct(self):
@@ -179,12 +191,14 @@ def deconstruct(self):
                 # Bind the deconstruct() to the instance to get the elidable
                 operation.deconstruct = deconstruct.__get__(operation, operation.__class__)
                 if not utils.is_code_in_site_packages(operation.code.__module__):
-                    code_name = unique_names.function(operation.code)
+                    code_name = utils.normalize_function_name(unique_names.function(operation.code))
                     operation.code = utils.copy_func(operation.code, code_name)
                     operation.code.__in_migration_file__ = True
                 if operation.reverse_code:
                     if not utils.is_code_in_site_packages(operation.reverse_code.__module__):
-                        reversed_code_name = unique_names.function(operation.reverse_code)
+                        reversed_code_name = utils.normalize_function_name(
+                            unique_names.function(operation.reverse_code)
+                        )
                         operation.reverse_code = utils.copy_func(operation.reverse_code, reversed_code_name)
                         operation.reverse_code.__in_migration_file__ = True
             elif isinstance(operation, dj_migrations.RunSQL):

django_squash/settings.py

@@ -2,3 +2,4 @@
 
 DJANGO_SQUASH_IGNORE_APPS = getattr(global_settings, "DJANGO_SQUASH_IGNORE_APPS", None) or []
 DJANGO_SQUASH_MIGRATION_NAME = getattr(global_settings, "DJANGO_SQUASH_MIGRATION_NAME", None) or "squashed"
+DJANGO_SQUASH_CUSTOM_RENAME_FUNCTION = getattr(global_settings, "DJANGO_SQUASH_CUSTOM_RENAME_FUNCTION", None) or None

docs/settings.rst

@@ -1,10 +1,10 @@
 Configuration
-~~~~~~~~~~~~~~~~~~~~~~~~
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 
 The following settings are available in order to customize your experience.
 
 ``DJANGO_SQUASH_IGNORE_APPS``
---------------------------------
+----------------------------------------
 
 Default: ``[]`` (Empty list)
 
@@ -13,8 +13,19 @@ Example: (``["app1", "app2"]``)
 Hardcoded list of apps to always ignore, no matter what, the same as ``--ignore`` in the ``./manage.py squash_migrations`` command.
 
 ``DJANGO_SQUASH_MIGRATION_NAME``
---------------------------------
+----------------------------------------
 
 Default: ``"squashed"`` (string)
 
 The generated migration name when ``./manage.py squash_migrations`` command is ran.
+
+``DJANGO_SQUASH_CUSTOM_RENAME_FUNCTION``
+----------------------------------------
+
+Default: ``None`` (None)
+
+Example: "path.to.generator_function"
+
+Dot path to the function that will rename the functions found inside ``RunPython`` operations.
+
+Function needs to accept 2 arguments: ``name`` (``str``) and ``context`` (``dict``) and must return a string (``-> str``)

tests/conftest.py

@@ -14,6 +14,8 @@
 from django.test.utils import extend_sys_path
 from django.utils.module_loading import module_dir
 
+from django_squash.db.migrations.utils import get_custom_rename_function
+
 from . import utils
 
 
@@ -149,6 +151,17 @@ def _call_squash_migrations(*args, **kwargs):
     yield _call_squash_migrations
 
 
+@pytest.fixture(autouse=True)
+def clear_get_custom_rename_function_cache():
+    """
+    To ensure that this function doesn't get cached with the wrong value and breaks tests,
+    we always clear it before and after each test.
+    """
+    get_custom_rename_function.cache_clear()
+    yield
+    get_custom_rename_function.cache_clear()
+
+
 @pytest.fixture
 def clean_db(django_db_blocker):
     """

tests/test_migrations.py

@@ -185,6 +185,45 @@ class Migration(migrations.Migration):
     assert pretty_extract_piece(app_squash, "") == expected
 
 
+def custom_func_naming(original_name, context):
+    """
+    Used in test_squashing_elidable_migration_unique_name_formatting to format the function names
+    """
+    return f"{context['migration'].name}_{original_name}"
+
+
+@pytest.mark.temporary_migration_module(module="app.tests.migrations.elidable", app_label="app")
+def test_squashing_elidable_migration_unique_name_formatting(migration_app_dir, call_squash_migrations, monkeypatch):
+    """
+    Test that DJANGO_SQUASH_CUSTOM_RENAME_FUNCTION integration works properly
+    """
+    monkeypatch.setattr(
+        "django_squash.settings.DJANGO_SQUASH_CUSTOM_RENAME_FUNCTION", f"{__name__}.custom_func_naming"
+    )
+
+    class Person(models.Model):
+        name = models.CharField(max_length=10)
+        dob = models.DateField()
+
+        class Meta:
+            app_label = "app"
+
+    call_squash_migrations()
+
+    app_squash = load_migration_module(migration_app_dir / "0004_squashed.py")
+    source = pretty_extract_piece(app_squash, "")
+    assert source.count("f_0002_person_age_same_name(") == 1
+    assert source.count("code=f_0002_person_age_same_name,") == 1
+    assert source.count("f_0002_person_age_same_name_2(") == 1
+    assert source.count("code=f_0002_person_age_same_name_2,") == 1
+    assert source.count("f_0003_add_dob_same_name(") == 1
+    assert source.count("code=f_0003_add_dob_same_name,") == 1
+    assert source.count("f_0003_add_dob_create_admin_MUST_ALWAYS_EXIST(") == 1
+    assert source.count("code=f_0003_add_dob_create_admin_MUST_ALWAYS_EXIST,") == 1
+    assert source.count("f_0003_add_dob_rollback_admin_MUST_ALWAYS_EXIST(") == 1
+    assert source.count("code=f_0003_add_dob_rollback_admin_MUST_ALWAYS_EXIST,") == 1
+
+
 @pytest.mark.temporary_migration_module(module="app.tests.migrations.simple", app_label="app")
 @pytest.mark.temporary_migration_module2(module="app2.tests.migrations.foreign_key", app_label="app2")
 def test_squashing_migration_simple(migration_app_dir, migration_app2_dir, call_squash_migrations):

tests/test_utils.py

@@ -65,14 +65,14 @@ def test_is_code_in_site_packages():
 
 
 def test_unique_names():
-    names = utils.UniqueVariableName()
+    names = utils.UniqueVariableName({})
     assert names("var") == "var"
     assert names("var") == "var_2"
     assert names("var_2") == "var_2_2"
 
 
 def test_unique_function_names_errors():
-    names = utils.UniqueVariableName()
+    names = utils.UniqueVariableName({})
 
     with pytest.raises(ValueError):
         names.function("not-a-function")
@@ -96,9 +96,27 @@ def test_unique_function_names_errors():
         names.function(D.func)
 
 
+def test_unique_function_names_context():
+    def custom_name(name, context):
+        return "{module}_{i}_{name}".format(**context, name=name)
+
+    names = utils.UniqueVariableName({"module": __name__.replace(".", "_")}, naming_function=custom_name)
+    collector = []
+    for i, func in enumerate((func2, func2_3, func2_impostor, func2_impostor2)):
+        names.update_context({"func": func, "i": i})
+        collector.append(names.function(func))
+
+    assert collector == [
+        "tests_test_utils_0_func2",
+        "tests_test_utils_1_func2_3",
+        "tests_test_utils_2_func2",
+        "tests_test_utils_3_func2",
+    ]
+
+
 def test_unique_function_names():
-    uniq1 = utils.UniqueVariableName()
-    uniq2 = utils.UniqueVariableName()
+    uniq1 = utils.UniqueVariableName({})
+    uniq2 = utils.UniqueVariableName({})
 
     reassigned_func2 = func2
     reassigned_func2_impostor = func2_impostor
@@ -150,3 +168,24 @@ def test_normalize_function_name():
     assert utils.normalize_function_name(reassigned_func2_impostor.__qualname__) == "func2"
     assert utils.normalize_function_name(A().func.__qualname__) == "func"
     assert utils.normalize_function_name(D.func.__qualname__) == "func"
+
+
+def test_get_custom_rename_function(monkeypatch):
+    """
+    Cover all cases where DJANGO_SQUASH_CUSTOM_RENAME_FUNCTION can go wrong
+    """
+    assert not utils.get_custom_rename_function()
+    utils.get_custom_rename_function.cache_clear()
+
+    monkeypatch.setattr("django_squash.settings.DJANGO_SQUASH_CUSTOM_RENAME_FUNCTION", "tests.test_utils.func2")
+    assert utils.get_custom_rename_function() == func2
+    utils.get_custom_rename_function.cache_clear()
+
+    monkeypatch.setattr("django_squash.settings.DJANGO_SQUASH_CUSTOM_RENAME_FUNCTION", "tests.test_utils.bad")
+    with pytest.raises(ImportError):
+        utils.get_custom_rename_function()
+    utils.get_custom_rename_function.cache_clear()
+
+    monkeypatch.setattr("django_squash.settings.DJANGO_SQUASH_CUSTOM_RENAME_FUNCTION", "does.not.exist")
+    with pytest.raises(ModuleNotFoundError):
+        utils.get_custom_rename_function()