New Protocol for migrations to make ty happy

sqlite_utils/migrations.py

@@ -1,19 +1,28 @@
 import datetime
 from collections.abc import Callable, Iterable
 from dataclasses import dataclass
-from typing import TYPE_CHECKING, cast
+from typing import TYPE_CHECKING, Protocol, TypeVar, cast
 
 if TYPE_CHECKING:
     from sqlite_utils.db import Database, Table
 
 
+class _MigrationFunction(Protocol):
+    __name__: str
+
+    def __call__(self, db: "Database", /) -> None: ...
+
+
+_MigrationFunctionT = TypeVar("_MigrationFunctionT", bound=_MigrationFunction)
+
+
 class Migrations:
     migrations_table = "_sqlite_migrations"
 
     @dataclass
     class _Migration:
         name: str
-        fn: Callable
+        fn: _MigrationFunction
         transactional: bool = True
 
     @dataclass
@@ -32,7 +41,7 @@ class Migrations:
 
     def __call__(
         self, *, name: str | None = None, transactional: bool = True
-    ) -> Callable:
+    ) -> Callable[[_MigrationFunctionT], _MigrationFunctionT]:
         """
         :param name: The name to use for this migration - if not provided,
           the name of the function will be used.
@@ -43,7 +52,7 @@ class Migrations:
           example those that execute ``VACUUM``.
         """
 
-        def inner(func: Callable) -> Callable:
+        def inner(func: _MigrationFunctionT) -> _MigrationFunctionT:
             migration_name = name or func.__name__
             if any(m.name == migration_name for m in self._migrations):
                 raise ValueError(