Coverage for python / lsst / scarlet / lite / io / migration.py: 50%

1from __future__ import annotations 

2 

3from typing import Callable, ClassVar 

4 

5__all__ = ["MigrationError", "MigrationRegistry", "migration"] 

6 

7Migrator = Callable[[dict], dict] 

8 

9PRE_SCHEMA = "0.0.0" 

10 

11 

12class MigrationError(Exception): 

13 """Custom error for migration issues.""" 

14 

15 pass 

16 

17 

18class MigrationRegistry: 

19 """Manages migration of data between different schema versions.""" 

20 

21 registry: ClassVar[dict[tuple[str, str], Migrator]] = {} 

22 current: ClassVar[dict[str, str]] = {} 

23 

24 @staticmethod 

25 def register(type_name: str, from_version: str, migrator: Migrator) -> None: 

26 """Register a migration function from one version to another.""" 

27 MigrationRegistry.registry[(type_name, from_version)] = migrator 

28 

29 @staticmethod 

30 def set_current(data_type: str, version: str) -> None: 

31 """Set the current version for a given data type.""" 

32 MigrationRegistry.current[data_type] = version 

33 

34 @staticmethod 

35 def migrate(data_type: str, data: dict) -> dict: 

36 """Migrate data to the current schema version. 

37 

38 Parameters 

39 ---------- 

40 data : 

41 The data to migrate. Must contain 'type' and 'version' keys. 

42 

43 Returns 

44 ------- 

45 result : 

46 The migrated data. 

47 """ 

48 if "version" not in data: 

49 # Unversioned data is pre-schema and is considered to be 

50 # version "0.0.0" for backwards compatibility. 

51 data["version"] = PRE_SCHEMA 

52 

53 from_version = data["version"] 

54 

55 if data_type not in MigrationRegistry.current: 

56 raise ValueError(f"No current version set for data type '{data_type}'.") 

57 

58 to_version = MigrationRegistry.current[data_type] 

59 

60 # Keep track of seen versions to avoid infinite loops 

61 seen: set[tuple[str, str]] = set() 

62 

63 while from_version != to_version: 

64 key = (data_type, from_version) 

65 if key not in MigrationRegistry.registry or key in seen: 

66 raise MigrationError( 

67 f"No migration path from version '{from_version}' for type '{data_type}'." 

68 ) 

69 

70 migrator = MigrationRegistry.registry[key] 

71 data = migrator(data) 

72 from_version = data["version"] 

73 

74 return data 

75 

76 

77def migration(type_name: str, from_version: str) -> Callable[[Migrator], Migrator]: 

78 """Decorator to register a migration step. 

79 

80 Parameters 

81 ---------- 

82 type_name : 

83 The type of data being migrated. 

84 from_version : 

85 The version the migrator converts from. 

86 

87 Returns 

88 ------- 

89 result : 

90 The decorator that registers the migration function. 

91 """ 

92 

93 def decorator(func: Migrator) -> Migrator: 

94 MigrationRegistry.register(type_name, from_version, func) 

95 return func 

96 

97 return decorator