Coverage for src/dataknobs_fsm/core/data_modes.py: 47%

1"""Data mode system for FSM state management.

2

3 This module defines the data handling strategies for state instances:

4- COPY: Deep copy on state entry, isolated modifications

5- REFERENCE: Lazy loading with optimistic locking

6- DIRECT: In-place modifications (single-threaded only)

7"""

8

9 from abc import ABC, abstractmethod

10 from copy import deepcopy

11 from enum import Enum

12 from threading import Lock

13 from typing import Any, Dict, TypeVar

14

15 T = TypeVar("T")

16

17

18 class DataHandlingMode(Enum):

19 """Data handling modes for state instances - defines how data is managed within states."""

20

21 COPY = "copy" # Deep copy on entry, isolated changes

22 REFERENCE = "reference" # Lazy loading with locking

23 DIRECT = "direct" # In-place modifications

24

25

26 class DataHandler(ABC):

27 """Abstract base class for data mode handlers."""

28

29 def __init__(self, mode: DataHandlingMode):

30 """Initialize the data handler.

31

32 Args:

33 mode: The data mode this handler implements.

34 """

35 self.mode = mode

36

37 @abstractmethod

38 def on_entry(self, data: Any) -> Any:

39 """Handle data when entering a state.

40

41 Args:

42 data: The input data to the state.

43

44 Returns:

45 The data to be used within the state.

46 """

47 pass

48

49 @abstractmethod

50 def on_modification(self, data: Any) -> Any:

51 """Handle data modification within a state.

52

53 Args:

54 data: The data being modified.

55

56 Returns:

57 The modified data.

58 """

59 pass

60

61 @abstractmethod

62 def on_exit(self, data: Any, commit: bool = True) -> Any:

63 """Handle data when exiting a state.

64

65 Args:

66 data: The state's data.

67 commit: Whether to commit changes.

68

69 Returns:

70 The final data after exit processing.

71 """

72 pass

73

74 @abstractmethod

75 def supports_concurrent_access(self) -> bool:

76 """Check if this handler supports concurrent access.

77

78 Returns:

79 True if concurrent access is supported.

80 """

81 pass

82

83

84 class CopyModeHandler(DataHandler):

85 """Handler for COPY mode - deep copy with isolated modifications."""

86

87 def __init__(self):

88 """Initialize the COPY mode handler."""

89 super().__init__(DataHandlingMode.COPY)

90 self._originals: Dict[int, Any] = {}

91 self._lock = Lock()

92

93 def on_entry(self, data: Any) -> Any:

94 """Create a deep copy of the data on state entry.

95

96 Args:

97 data: The input data to copy.

98

99 Returns:

100 A deep copy of the data.

101 """

102 copied = deepcopy(data)

103 with self._lock:

104 self._originals[id(copied)] = data

105 return copied

106

107 def on_modification(self, data: Any) -> Any:

108 """Allow modifications to the copied data.

109

110 Args:

111 data: The copied data being modified.

112

113 Returns:

114 The same data (modifications are isolated).

115 """

116 return data

117

118 def on_exit(self, data: Any, commit: bool = True) -> Any:

119 """Handle exit, optionally committing changes.

120

121 Args:

122 data: The modified copy.

123 commit: Whether to return modified or original data.

124

125 Returns:

126 Modified data if commit=True, original otherwise.

127 """

128 with self._lock:

129 original = self._originals.pop(id(data), None)

130

131 if commit:

132 return data

133 else:

134 return original if original is not None else data

135

136 def supports_concurrent_access(self) -> bool:

137 """COPY mode supports concurrent access.

138

139 Returns:

140 True, as each state gets its own copy.

141 """

142 return True

143

144

145 class ReferenceModeHandler(DataHandler):

146 """Handler for REFERENCE mode - lazy loading with optimistic locking."""

147

148 def __init__(self):

149 """Initialize the REFERENCE mode handler."""

150 super().__init__(DataHandlingMode.REFERENCE)

151 # Don't use WeakValueDictionary as it can't hold dict objects

152 self._references: Dict[int, Any] = {}

153 self._versions: Dict[int, int] = {}

154 self._locks: Dict[int, Lock] = {}

155 self._global_lock = Lock()

156

157 def on_entry(self, data: Any) -> Any:

158 """Store a reference to the data with versioning.

159

160 Args:

161 data: The input data to reference.

162

163 Returns:

164 The same data object (by reference).

165 """

166 data_id = id(data)

167 with self._global_lock:

168 if data_id not in self._references:

169 self._references[data_id] = data

170 self._versions[data_id] = 0

171 self._locks[data_id] = Lock()

172 return data

173

174 def on_modification(self, data: Any) -> Any:

175 """Track modifications with optimistic locking.

176

177 Args:

178 data: The data being modified.

179

180 Returns:

181 The same data object.

182 """

183 data_id = id(data)

184 with self._global_lock:

185 if data_id in self._locks:

186 lock = self._locks[data_id]

187 else:

188 lock = Lock()

189 self._locks[data_id] = lock

190

191 with lock:

192 if data_id in self._versions:

193 self._versions[data_id] += 1

194

195 return data

196

197 def on_exit(self, data: Any, commit: bool = True) -> Any:

198 """Clean up references on exit.

199

200 Args:

201 data: The referenced data.

202 commit: Whether changes should be kept.

203

204 Returns:

205 The data object.

206 """

207 data_id = id(data)

208

209 if not commit:

210 # In REFERENCE mode, we can't truly rollback in-place changes

211 # This would require snapshot/restore functionality

212 pass

213

214 # Clean up tracking if no longer needed

215 with self._global_lock:

216 # Clean up if version is 0 (no modifications)

217 if data_id in self._versions and self._versions[data_id] == 0:

218 self._references.pop(data_id, None)

219 self._versions.pop(data_id, None)

220 self._locks.pop(data_id, None)

221

222 return data

223

224 def supports_concurrent_access(self) -> bool:

225 """REFERENCE mode supports concurrent access with locking.

226

227 Returns:

228 True, with optimistic locking for safety.

229 """

230 return True

231

232

233 class DirectModeHandler(DataHandler):

234 """Handler for DIRECT mode - in-place modifications without copying."""

235

236 def __init__(self):

237 """Initialize the DIRECT mode handler."""

238 super().__init__(DataHandlingMode.DIRECT)

239 self._active_data: Any | None = None

240 self._lock = Lock()

241

242 def on_entry(self, data: Any) -> Any:

243 """Use data directly without copying.

244

245 Args:

246 data: The input data to use directly.

247

248 Returns:

249 The same data object.

250 """

251 with self._lock:

252 if self._active_data is not None:

253 raise RuntimeError(

254 "DIRECT mode does not support concurrent access. "

255 "Another state is currently using direct mode."

256 )

257 self._active_data = data

258 return data

259

260 def on_modification(self, data: Any) -> Any:

261 """Allow direct in-place modifications.

262

263 Args:

264 data: The data being modified.

265

266 Returns:

267 The same data object.

268 """

269 return data

270

271 def on_exit(self, data: Any, commit: bool = True) -> Any:

272 """Clear the active data reference.

273

274 Args:

275 data: The data object.

276 commit: Whether to keep changes (always True for DIRECT).

277

278 Returns:

279 The data object.

280 """

281 with self._lock:

282 self._active_data = None

283

284 # In DIRECT mode, changes are always committed (in-place)

285 return data

286

287 def supports_concurrent_access(self) -> bool:

288 """DIRECT mode does not support concurrent access.

289

290 Returns:

291 False, as only one state can use DIRECT mode at a time.

292 """

293 return False

294

295

296 class DataModeManager:

297 """Manager for data mode handlers."""

298

299 def __init__(self, default_mode: DataHandlingMode = DataHandlingMode.COPY):

300 """Initialize the data mode manager.

301

302 Args:

303 default_mode: The default data mode to use.

304 """

305 self.default_mode = default_mode

306 self._handlers: Dict[DataHandlingMode, DataHandler] = {

307 DataHandlingMode.COPY: CopyModeHandler(),

308 DataHandlingMode.REFERENCE: ReferenceModeHandler(),

309 DataHandlingMode.DIRECT: DirectModeHandler(),

310 }

311

312 def get_handler(self, mode: DataHandlingMode | None = None) -> DataHandler:

313 """Get a handler for the specified mode.

314

315 Args:

316 mode: The data mode, or None for default.

317

318 Returns:

319 The appropriate data handler.

320 """

321 if mode is None:

322 mode = self.default_mode

323 return self._handlers[mode]

324

325 def set_default_mode(self, mode: DataHandlingMode) -> None:

326 """Set the default data mode.

327

328 Args:

329 mode: The new default data mode.

330 """

331 self.default_mode = mode

332

333

334# Global registry of data handlers

335 _GLOBAL_HANDLERS = {

336 DataHandlingMode.COPY: CopyModeHandler(),

337 DataHandlingMode.REFERENCE: ReferenceModeHandler(),

338 DataHandlingMode.DIRECT: DirectModeHandler(),

339}

340

341

342 def get_data_handler(mode: DataHandlingMode) -> DataHandler:

343 """Get a data handler for the specified mode.

344

345 Args:

346 mode: The data mode.

347

348 Returns:

349 The appropriate data handler.

350 """

351 return _GLOBAL_HANDLERS[mode]