Add typealias for migration classes

Author: gmazoyer

backend/infrahub/cli/db.py

@@ -63,12 +63,7 @@ def get_timestamp_string() -> str:
 
 if TYPE_CHECKING:
     from infrahub.cli.context import CliContext
-    from infrahub.core.migrations.shared import (
-        ArbitraryMigration,
-        GraphMigration,
-        InternalSchemaMigration,
-        MigrationRequiringRebase,
-    )
+    from infrahub.core.migrations.shared import MigrationTypes
     from infrahub.database import InfrahubDatabase
     from infrahub.database.index import IndexManagerBase
 
@@ -287,10 +282,10 @@ async def index(
 
 async def detect_migration_to_run(
     current_graph_version: int, migration_number: int | str | None = None
-) -> Sequence[GraphMigration | InternalSchemaMigration | ArbitraryMigration | MigrationRequiringRebase]:
+) -> Sequence[MigrationTypes]:
     """Return a sequence of migrations to apply to upgrade the database."""
     rprint("Checking current state of the database")
-    migrations: list[GraphMigration | InternalSchemaMigration | ArbitraryMigration | MigrationRequiringRebase] = []
+    migrations: list[MigrationTypes] = []
 
     if migration_number:
         migration = get_migration_by_number(migration_number)
@@ -316,9 +311,7 @@ async def detect_migration_to_run(
 
 
 async def migrate_database(
-    db: InfrahubDatabase,
-    migrations: Sequence[GraphMigration | InternalSchemaMigration | ArbitraryMigration | MigrationRequiringRebase],
-    initialize: bool = False,
+    db: InfrahubDatabase, migrations: Sequence[MigrationTypes], initialize: bool = False
 ) -> bool:
     """Apply the latest migrations to the database, this function will print the status directly in the console.
 

backend/infrahub/core/migrations/graph/__init__.py

@@ -47,9 +47,10 @@
 from .m043_backfill_hfid_display_label_in_db import Migration043
 
 if TYPE_CHECKING:
-    from ..shared import ArbitraryMigration, GraphMigration, InternalSchemaMigration, MigrationRequiringRebase
+    from ..shared import MigrationTypes
 
-MIGRATIONS: list[type[GraphMigration | InternalSchemaMigration | ArbitraryMigration | MigrationRequiringRebase]] = [
+
+MIGRATIONS: list[type[MigrationTypes]] = [
     Migration001,
     Migration002,
     Migration003,
@@ -96,9 +97,7 @@
 ]
 
 
-async def get_graph_migrations(
-    current_graph_version: int,
-) -> Sequence[GraphMigration | InternalSchemaMigration | ArbitraryMigration | MigrationRequiringRebase]:
+async def get_graph_migrations(current_graph_version: int) -> Sequence[MigrationTypes]:
     applicable_migrations = []
     for migration_class in MIGRATIONS:
         migration = migration_class.init()
@@ -109,9 +108,7 @@ async def get_graph_migrations(
     return applicable_migrations
 
 
-def get_migration_by_number(
-    migration_number: int | str,
-) -> GraphMigration | InternalSchemaMigration | ArbitraryMigration | MigrationRequiringRebase:
+def get_migration_by_number(migration_number: int | str) -> MigrationTypes:
     # Convert to string and pad with zeros if needed
     try:
         num = int(migration_number)

backend/infrahub/core/migrations/shared.py

@@ -1,6 +1,6 @@
 from __future__ import annotations
 
-from typing import TYPE_CHECKING, Any, Sequence
+from typing import TYPE_CHECKING, Any, Sequence, TypeAlias
 
 from pydantic import BaseModel, ConfigDict, Field
 from typing_extensions import Self
@@ -245,3 +245,6 @@ async def execute_against_branch(self, db: InfrahubDatabase, branch: Branch) ->
     async def execute(self, db: InfrahubDatabase) -> MigrationResult:
         """Method that will be run against the default branch."""
         raise NotImplementedError()
+
+
+MigrationTypes: TypeAlias = GraphMigration | InternalSchemaMigration | ArbitraryMigration | MigrationRequiringRebase