Coverage for src/dataknobs_fsm/io/base.py: 86%

1"""Base I/O abstraction layer components.

2

3 This module provides the base abstractions for unified I/O operations.

4"""

5

6 from abc import ABC, abstractmethod

7 from dataclasses import dataclass, field

8 from enum import Enum

9 from typing import (

10 Any, Dict, List, Union, AsyncIterator, Iterator,

11 Callable, TypeVar, Protocol

12)

13

14 T = TypeVar('T')

15

16

17 class IOMode(Enum):

18 """I/O operation modes."""

19 READ = "read"

20 WRITE = "write"

21 APPEND = "append"

22 STREAM = "stream"

23 BATCH = "batch"

24

25

26 class IOFormat(Enum):

27 """Supported I/O formats."""

28 JSON = "json"

29 CSV = "csv"

30 XML = "xml"

31 PARQUET = "parquet"

32 AVRO = "avro"

33 TEXT = "text"

34 BINARY = "binary"

35 DATABASE = "database"

36 API = "api"

37

38

39@dataclass

40 class IOConfig:

41 """Configuration for I/O operations."""

42 mode: IOMode

43 format: IOFormat

44 source: Union[str, Dict[str, Any]] # Path, URL, or connection config

45 target: Union[str, Dict[str, Any]] | None = None

46

47 # Operation settings

48 batch_size: int = 1000

49 buffer_size: int = 8192

50 timeout: float | None = None

51 retry_count: int = 3

52 retry_delay: float = 1.0

53

54 # Format-specific settings

55 encoding: str = "utf-8"

56 compression: str | None = None

57 delimiter: str = ","

58 headers: Dict[str, str] | None = None

59

60 # Advanced settings

61 parallel_workers: int = 1

62 checkpoint_enabled: bool = False

63 checkpoint_interval: int = 10000

64 error_handler: Callable[[Exception, Any], Any] | None = None

65

66 # Additional options

67 options: Dict[str, Any] = field(default_factory=dict)

68

69

70 class IOProvider(ABC):

71 """Base I/O provider interface."""

72

73 def __init__(self, config: IOConfig):

74 """Initialize provider with configuration."""

75 self.config = config

76 self._is_open = False

77

78 @abstractmethod

79 def open(self) -> None:

80 """Open the I/O connection."""

81 pass

82

83 @abstractmethod

84 def close(self) -> None:

85 """Close the I/O connection."""

86 pass

87

88 @abstractmethod

89 def validate(self) -> bool:

90 """Validate the I/O configuration and connection."""

91 pass

92

93 @property

94 def is_open(self) -> bool:

95 """Check if connection is open."""

96 return self._is_open

97

98 def __enter__(self):

99 """Context manager entry."""

100 self.open()

101 return self

102

103 def __exit__(self, exc_type, exc_val, exc_tb):

104 """Context manager exit."""

105 self.close()

106

107

108 class AsyncIOProvider(IOProvider):

109 """Async I/O provider interface."""

110

111 @abstractmethod

112 async def read(self, **kwargs) -> Any:

113 """Read data asynchronously."""

114 pass

115

116 @abstractmethod

117 async def write(self, data: Any, **kwargs) -> None:

118 """Write data asynchronously."""

119 pass

120

121 @abstractmethod

122 async def stream_read(self, **kwargs) -> AsyncIterator[Any]:

123 """Stream read data asynchronously."""

124 pass

125

126 @abstractmethod

127 async def stream_write(self, data_stream: AsyncIterator[Any], **kwargs) -> None:

128 """Stream write data asynchronously."""

129 pass

130

131 @abstractmethod

132 async def batch_read(self, batch_size: int | None = None, **kwargs) -> AsyncIterator[List[Any]]:

133 """Read data in batches asynchronously."""

134 pass

135

136 @abstractmethod

137 async def batch_write(self, batches: AsyncIterator[List[Any]], **kwargs) -> None:

138 """Write data in batches asynchronously."""

139 pass

140

141 async def open(self) -> None:

142 """Open the async I/O connection."""

143 self._is_open = True

144

145 async def close(self) -> None:

146 """Close the async I/O connection."""

147 self._is_open = False

148

149 async def __aenter__(self):

150 """Async context manager entry."""

151 await self.open()

152 return self

153

154 async def __aexit__(self, exc_type, exc_val, exc_tb):

155 """Async context manager exit."""

156 await self.close()

157

158

159 class SyncIOProvider(IOProvider):

160 """Synchronous I/O provider interface."""

161

162 @abstractmethod

163 def read(self, **kwargs) -> Any:

164 """Read data synchronously."""

165 pass

166

167 @abstractmethod

168 def write(self, data: Any, **kwargs) -> None:

169 """Write data synchronously."""

170 pass

171

172 @abstractmethod

173 def stream_read(self, **kwargs) -> Iterator[Any]:

174 """Stream read data synchronously."""

175 pass

176

177 @abstractmethod

178 def stream_write(self, data_stream: Iterator[Any], **kwargs) -> None:

179 """Stream write data synchronously."""

180 pass

181

182 @abstractmethod

183 def batch_read(self, batch_size: int | None = None, **kwargs) -> Iterator[List[Any]]:

184 """Read data in batches synchronously."""

185 pass

186

187 @abstractmethod

188 def batch_write(self, batches: Iterator[List[Any]], **kwargs) -> None:

189 """Write data in batches synchronously."""

190 pass

191

192 def open(self) -> None:

193 """Open the sync I/O connection."""

194 self._is_open = True

195

196 def close(self) -> None:

197 """Close the sync I/O connection."""

198 self._is_open = False

199

200

201 class TransformProtocol(Protocol):

202 """Protocol for data transformations."""

203

204 def transform(self, data: Any) -> Any:

205 """Transform data."""

206 ...

207

208 async def async_transform(self, data: Any) -> Any:

209 """Transform data asynchronously."""

210 ...

211

212

213 class IOAdapter(ABC):

214 """Base adapter for converting between different I/O providers."""

215

216 @abstractmethod

217 def adapt_config(self, config: IOConfig) -> Dict[str, Any]:

218 """Adapt configuration for specific provider."""

219 pass

220

221 @abstractmethod

222 def adapt_data(self, data: Any, direction: IOMode) -> Any:

223 """Adapt data format for specific provider."""

224 pass

225

226 @abstractmethod

227 def create_provider(self, config: IOConfig, is_async: bool = True) -> IOProvider:

228 """Create appropriate provider instance."""

229 pass