Coverage for src/dataknobs_fsm/functions/base.py: 53%

1"""Base interfaces and classes for FSM functions.

2

3 This module defines the interfaces for:

4- Validation functions (check data validity)

5- Transform functions (modify data)

6- State test functions (determine next state)

7- End state test functions (check if processing should end)

8- Resources (external systems and services)

9"""

10

11 from abc import ABC, abstractmethod

12 from typing import Any, Dict, List, Tuple, TypeVar

13 from enum import Enum

14 from dataclasses import dataclass, field

15

16 T = TypeVar("T")

17

18

19 class FunctionType(Enum):

20 """Types of functions in the FSM."""

21

22 VALIDATION = "validation"

23 TRANSFORM = "transform"

24 STATE_TEST = "state_test"

25 END_STATE_TEST = "end_state_test"

26

27

28 class ExecutionResult:

29 """Result of function execution."""

30

31 def __init__(

32 self,

33 success: bool,

34 data: Any | None = None,

35 error: str | None = None,

36 metadata: Dict[str, Any] | None = None

37 ):

38 """Initialize execution result.

39

40 Args:

41 success: Whether execution succeeded.

42 data: Result data if successful.

43 error: Error message if failed.

44 metadata: Additional metadata about execution.

45 """

46 self.success = success

47 self.data = data

48 self.error = error

49 self.metadata = metadata or {}

50

51 @classmethod

52 def success_result(cls, data: Any, metadata: Dict[str, Any] | None = None) -> 'ExecutionResult':

53 """Create a successful result.

54

55 Args:

56 data: The result data.

57 metadata: Optional metadata.

58

59 Returns:

60 A successful ExecutionResult.

61 """

62 return cls(success=True, data=data, metadata=metadata)

63

64 @classmethod

65 def failure_result(cls, error: str, metadata: Dict[str, Any] | None = None) -> 'ExecutionResult':

66 """Create a failure result.

67

68 Args:

69 error: The error message.

70 metadata: Optional metadata.

71

72 Returns:

73 A failed ExecutionResult.

74 """

75 return cls(success=False, error=error, metadata=metadata)

76

77 def to_dict(self) -> Dict[str, Any]:

78 """Convert to dictionary for serialization.

79

80 Returns:

81 Dictionary representation of the result.

82 """

83 return {

84 'success': self.success,

85 'data': self.data,

86 'error': self.error,

87 'metadata': self.metadata

88 }

89

90 def __json__(self) -> Dict[str, Any]:

91 """Support JSON serialization.

92

93 Returns:

94 Dictionary representation for JSON.

95 """

96 return self.to_dict()

97

98

99@dataclass

100 class FunctionContext:

101 """Context passed to functions during execution."""

102 state_name: str

103 function_name: str

104 metadata: Dict[str, Any] = field(default_factory=dict)

105 resources: Dict[str, Any] = field(default_factory=dict)

106 variables: Dict[str, Any] = field(default_factory=dict) # Shared variables

107 network_name: str | None = None # Current network for scoping

108

109

110 class IValidationFunction(ABC):

111 """Interface for validation functions."""

112

113 @abstractmethod

114 def validate(self, data: Any, context: Dict[str, Any] | None = None) -> ExecutionResult:

115 """Validate data according to function logic.

116

117 Args:

118 data: The data to validate.

119 context: Optional execution context.

120

121 Returns:

122 ExecutionResult with validation outcome.

123 """

124 pass

125

126 @abstractmethod

127 def get_validation_rules(self) -> Dict[str, Any]:

128 """Get the validation rules this function implements.

129

130 Returns:

131 Dictionary describing the validation rules.

132 """

133 pass

134

135

136 class ITransformFunction(ABC):

137 """Interface for transform functions."""

138

139 @abstractmethod

140 def transform(self, data: Any, context: Dict[str, Any] | None = None) -> ExecutionResult:

141 """Transform data according to function logic.

142

143 Args:

144 data: The data to transform.

145 context: Optional execution context.

146

147 Returns:

148 ExecutionResult with transformed data.

149 """

150 pass

151

152 @abstractmethod

153 def get_transform_description(self) -> str:

154 """Get a description of the transformation.

155

156 Returns:

157 String describing what this transform does.

158 """

159 pass

160

161

162 class IStateTestFunction(ABC):

163 """Interface for state test functions."""

164

165 @abstractmethod

166 def test(self, data: Any, context: Dict[str, Any] | None = None) -> Tuple[bool, str | None]:

167 """Test if a condition is met for state transition.

168

169 Args:

170 data: The data to test.

171 context: Optional execution context.

172

173 Returns:

174 Tuple of (test_passed, reason).

175 """

176 pass

177

178 @abstractmethod

179 def get_test_description(self) -> str:

180 """Get a description of what this test checks.

181

182 Returns:

183 String describing the test condition.

184 """

185 pass

186

187

188 class IEndStateTestFunction(ABC):

189 """Interface for end state test functions."""

190

191 @abstractmethod

192 def should_end(self, data: Any, context: Dict[str, Any] | None = None) -> Tuple[bool, str | None]:

193 """Test if processing should end.

194

195 Args:

196 data: The current data.

197 context: Optional execution context.

198

199 Returns:

200 Tuple of (should_end, reason).

201 """

202 pass

203

204 @abstractmethod

205 def get_end_condition(self) -> str:

206 """Get a description of the end condition.

207

208 Returns:

209 String describing when processing ends.

210 """

211 pass

212

213

214 class ResourceStatus(Enum):

215 """Status of a resource."""

216

217 UNINITIALIZED = "uninitialized"

218 INITIALIZING = "initializing"

219 READY = "ready"

220 BUSY = "busy"

221 ERROR = "error"

222 SHUTDOWN = "shutdown"

223

224

225@dataclass

226 class ResourceConfig:

227 """Configuration for a resource."""

228

229 name: str

230 type: str

231 connection_params: Dict[str, Any]

232 pool_size: int | None = None

233 timeout: float | None = None

234 retry_policy: Dict[str, Any] | None = None

235 health_check_interval: float | None = None

236

237

238 class IResource(ABC):

239 """Interface for external resources."""

240

241 @abstractmethod

242 async def initialize(self, config: ResourceConfig) -> None:

243 """Initialize the resource.

244

245 Args:

246 config: Resource configuration.

247 """

248 pass

249

250 @abstractmethod

251 async def acquire(self, timeout: float | None = None) -> Any:

252 """Acquire a connection/handle to the resource.

253

254 Args:

255 timeout: Optional timeout for acquisition.

256

257 Returns:

258 A resource handle/connection.

259 """

260 pass

261

262 @abstractmethod

263 async def release(self, handle: Any) -> None:

264 """Release a resource handle/connection.

265

266 Args:

267 handle: The handle to release.

268 """

269 pass

270

271 @abstractmethod

272 async def health_check(self) -> bool:

273 """Check if the resource is healthy.

274

275 Returns:

276 True if healthy, False otherwise.

277 """

278 pass

279

280 @abstractmethod

281 async def shutdown(self) -> None:

282 """Shutdown the resource and cleanup."""

283 pass

284

285 @abstractmethod

286 def get_status(self) -> ResourceStatus:

287 """Get the current resource status.

288

289 Returns:

290 Current ResourceStatus.

291 """

292 pass

293

294

295# Exception classes

296

297 class FSMError(Exception):

298 """Base exception for FSM errors."""

299 pass

300

301

302 class ValidationError(FSMError):

303 """Raised when validation fails."""

304

305 def __init__(self, message: str, validation_errors: List[str] | None = None):

306 """Initialize validation error.

307

308 Args:

309 message: Error message.

310 validation_errors: List of specific validation errors.

311 """

312 super().__init__(message)

313 self.validation_errors = validation_errors or []

314

315

316 class TransformError(FSMError):

317 """Raised when transformation fails."""

318 pass

319

320

321 class StateTransitionError(FSMError):

322 """Raised when state transition fails."""

323

324 def __init__(self, message: str, from_state: str, to_state: str | None = None):

325 """Initialize state transition error.

326

327 Args:

328 message: Error message.

329 from_state: The state transitioning from.

330 to_state: The state attempting to transition to.

331 """

332 super().__init__(message)

333 self.from_state = from_state

334 self.to_state = to_state

335

336

337 class ResourceError(FSMError):

338 """Raised when resource operations fail."""

339

340 def __init__(self, message: str, resource_name: str, operation: str):

341 """Initialize resource error.

342

343 Args:

344 message: Error message.

345 resource_name: Name of the resource.

346 operation: The operation that failed.

347 """

348 super().__init__(message)

349 self.resource_name = resource_name

350 self.operation = operation

351

352

353 class ConfigurationError(FSMError):

354 """Raised when configuration is invalid."""

355 pass

356

357

358# Base implementations

359

360 class BaseFunction:

361 """Base class for functions with common functionality."""

362

363 def __init__(self, name: str, description: str = ""):

364 """Initialize base function.

365

366 Args:

367 name: Function name.

368 description: Function description.

369 """

370 self.name = name

371 self.description = description

372 self.execution_count = 0

373 self.error_count = 0

374

375 def _record_execution(self, success: bool) -> None:

376 """Record execution statistics.

377

378 Args:

379 success: Whether execution succeeded.

380 """

381 self.execution_count += 1

382 if not success:

383 self.error_count += 1

384

385 def get_stats(self) -> Dict[str, int]:

386 """Get execution statistics.

387

388 Returns:

389 Dictionary with execution stats.

390 """

391 return {

392 "executions": self.execution_count,

393 "errors": self.error_count,

394 "success_rate": float( # type: ignore

395 (self.execution_count - self.error_count) / self.execution_count

396 if self.execution_count > 0 else 0

397 )

398 }

399

400

401 class CompositeFunction(BaseFunction):

402 """Base class for functions that compose multiple sub-functions."""

403

404 def __init__(self, name: str, functions: List[BaseFunction], description: str = ""):

405 """Initialize composite function.

406

407 Args:

408 name: Function name.

409 functions: List of sub-functions to compose.

410 description: Function description.

411 """

412 super().__init__(name, description)

413 self.functions = functions

414

415 def add_function(self, function: BaseFunction) -> None:

416 """Add a function to the composite.

417

418 Args:

419 function: Function to add.

420 """

421 self.functions.append(function)

422

423 def remove_function(self, function_name: str) -> bool:

424 """Remove a function from the composite.

425

426 Args:

427 function_name: Name of function to remove.

428

429 Returns:

430 True if removed, False if not found.

431 """

432 for i, func in enumerate(self.functions):

433 if func.name == function_name:

434 self.functions.pop(i)

435 return True

436 return False

437

438

439# Simple Function class for basic use

440 class Function(ABC):

441 """Abstract base class for simple functions."""

442

443 @abstractmethod

444 def execute(self, data: Any, context: 'FunctionContext') -> Any:

445 """Execute the function.

446

447 Args:

448 data: Input data.

449 context: Function context.

450

451 Returns:

452 Function result.

453 """

454 pass

455

456

457# FunctionRegistry for managing functions

458 class FunctionRegistry:

459 """Registry for managing FSM functions."""

460

461 def __init__(self):

462 """Initialize function registry."""

463 self.functions: Dict[str, Any] = {}

464 self.validators: Dict[str, IValidationFunction] = {}

465 self.transforms: Dict[str, ITransformFunction] = {}

466

467 def register(self, name: str, function: Any) -> None:

468 """Register a function.

469

470 Args:

471 name: Function name.

472 function: Function instance.

473 """

474 if isinstance(function, Function):

475 self.functions[name] = function

476 elif isinstance(function, IValidationFunction):

477 self.validators[name] = function

478 elif isinstance(function, ITransformFunction):

479 self.transforms[name] = function

480 else:

481 # Store as generic function

482 self.functions[name] = function

483

484 def get_function(self, name: str) -> Any | None:

485 """Get a function by name.

486

487 Args:

488 name: Function name.

489

490 Returns:

491 Function instance or None.

492 """

493 # Check all registries

494 if name in self.functions:

495 return self.functions[name]

496 elif name in self.validators:

497 return self.validators[name]

498 elif name in self.transforms:

499 return self.transforms[name]

500 return None

501

502 def remove(self, name: str) -> bool:

503 """Remove a function.

504

505 Args:

506 name: Function name.

507

508 Returns:

509 True if removed.

510 """

511 if name in self.functions:

512 del self.functions[name]

513 return True

514 elif name in self.validators:

515 del self.validators[name]

516 return True

517 elif name in self.transforms:

518 del self.transforms[name]

519 return True

520 return False

521

522 def list_functions(self) -> List[str]:

523 """List all registered functions.

524

525 Returns:

526 List of function names.

527 """

528 all_names = []

529 all_names.extend(self.functions.keys())

530 all_names.extend(self.validators.keys())

531 all_names.extend(self.transforms.keys())

532 return sorted(all_names)

533

534 def clear(self) -> None:

535 """Clear all registered functions."""

536 self.functions.clear()

537 self.validators.clear()

538 self.transforms.clear()

539

540

541# Alias FunctionError to StateTransitionError for compatibility

542 FunctionError = StateTransitionError