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

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

3This 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"""

11from abc import ABC, abstractmethod

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

13from enum import Enum

14from dataclasses import dataclass, field

16T = TypeVar("T")

19class FunctionType(Enum):

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

22 VALIDATION = "validation"

23 TRANSFORM = "transform"

24 STATE_TEST = "state_test"

25 END_STATE_TEST = "end_state_test"

28class ExecutionResult:

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

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.

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 {}

51 @classmethod

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

53 """Create a successful result.

55 Args:

56 data: The result data.

57 metadata: Optional metadata.

59 Returns:

60 A successful ExecutionResult.

61 """

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

64 @classmethod

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

66 """Create a failure result.

68 Args:

69 error: The error message.

70 metadata: Optional metadata.

72 Returns:

73 A failed ExecutionResult.

74 """

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

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

78 """Convert to dictionary for serialization.

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 }

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

91 """Support JSON serialization.

93 Returns:

94 Dictionary representation for JSON.

95 """

96 return self.to_dict()

99@dataclass

100class 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

110class 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

136class 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

162class 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

188class 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

214class 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

226class 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

238class 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

297class FSMError(Exception):

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

299 pass

300

301

302class 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

316class TransformError(FSMError):

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

318 pass

319

320

321class 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

337class 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

353class ConfigurationError(FSMError):

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

355 pass

356

357

358# Base implementations

359

360class 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

401class 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

440class 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

458class 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

542FunctionError = StateTransitionError