Coverage for jinja2_async_environment/loaders/base.py: 80%

1"""Base classes and protocols for async template loaders."""

2

3 import typing as t

4 from abc import abstractmethod

5

6 from anyio import Path as AsyncPath

7 from jinja2.loaders import BaseLoader

8 from jinja2.utils import internalcode

9

10 if t.TYPE_CHECKING:

11 from jinja2 import Template

12

13 from ..environment import AsyncEnvironment

14

15# Define TemplateData type for better type checking

16# Replaced namedtuple with TypedCache for better type safety

17

18

19 class TemplateDataType(t.NamedTuple):

20 """Type definition for template data."""

21

22 source: t.Any

23 path: t.Any

24 uptodate: t.Any

25 source_str: str

26 name: str

27

28

29# Type alias for source data

30 SourceType = tuple[str | bytes, str | None, t.Callable[[], bool] | None]

31

32

33 class AsyncLoaderProtocol(t.Protocol):

34 """Protocol defining the interface for async template loaders.

35

36 This protocol ensures all async loaders implement the required methods

37 for loading templates asynchronously while maintaining compatibility

38 with Jinja2's loader interface.

39 """

40

41 async def get_source_async(

42 self, environment: "AsyncEnvironment", name: str

43 ) -> SourceType | None:

44 """Get template source asynchronously.

45

46 Args:

47 environment: The async environment instance

48 name: Template name to load

49

50 Returns:

51 Tuple of (source, filename, uptodate_func) or None if not found

52

53 Raises:

54 TemplateNotFound: If template cannot be found

55 """

56 ...

57

58 async def list_templates_async(self) -> list[str]:

59 """List all available templates asynchronously.

60

61 Returns:

62 List of template names

63

64 Raises:

65 TypeError: If listing is not supported by this loader

66 """

67 ...

68

69 async def load_async(

70 self,

71 environment: "AsyncEnvironment",

72 name: str,

73 env_globals: dict[str, t.Any] | None = None,

74 ) -> "Template":

75 """Load template asynchronously.

76

77 Args:

78 environment: The async environment instance

79 name: Template name to load

80 env_globals: Global variables for the template

81

82 Returns:

83 Compiled Template object

84

85 Raises:

86 TemplateNotFound: If template cannot be found

87 """

88 ...

89

90

91 class AsyncBaseLoader(BaseLoader):

92 """Base class for async template loaders with memory optimization.

93

94 This class provides the foundation for all async loaders, implementing

95 common functionality while using __slots__ for memory efficiency.

96

97 Features:

98 - Thread-safe lazy initialization

99 - Memory-optimized with __slots__

100 - Robust error handling and resource management

101 - Compatible with Jinja2's loader interface

102 - Support for async template operations

103

104 Thread Safety:

105 The loader implements thread-safe initialization using RLock to ensure

106 that multiple concurrent calls to _ensure_initialized() are handled safely.

107 """

108

109 __slots__ = ("searchpath", "_initialized", "_init_lock")

110

111 searchpath: list[AsyncPath]

112 _initialized: bool

113 _init_lock: t.Any | None

114

115 def __init__(self, searchpath: AsyncPath | str | t.Sequence[AsyncPath | str]):

116 """Initialize the async loader.

117

118 Args:

119 searchpath: Path or sequence of paths to search for templates

120

121 Raises:

122 TypeError: If searchpath is not a valid type

123 ValueError: If searchpath is empty or contains invalid paths

124 """

125 # Validate and normalize searchpath

126 searchpath_list = self._normalize_searchpath(searchpath)

127

128 # Convert to AsyncPath objects for consistency and validate

129 self.searchpath = self._convert_to_async_paths(searchpath_list)

130

131 self._initialized = False

132 self._init_lock = None # Will be created on first use

133

134 def _normalize_searchpath(

135 self, searchpath: AsyncPath | str | t.Sequence[AsyncPath | str]

136 ) -> list[AsyncPath | str]:

137 """Normalize searchpath to a list of paths.

138

139 Args:

140 searchpath: Path or sequence of paths to normalize

141

142 Returns:

143 List of normalized paths

144

145 Raises:

146 TypeError: If searchpath is not a valid type

147 ValueError: If searchpath is empty or contains invalid paths

148 """

149 if isinstance(searchpath, str) or hasattr(

150 searchpath, "parts"

151 ): # AsyncPath check

152 return [t.cast(AsyncPath | str, searchpath)]

153 # Try to treat as sequence

154 return self._normalize_sequence_searchpath(searchpath)

155

156 def _normalize_sequence_searchpath(

157 self, searchpath: AsyncPath | str | t.Sequence[AsyncPath | str]

158 ) -> list[AsyncPath | str]:

159 """Normalize sequence searchpath to a list of paths.

160

161 Args:

162 searchpath: Sequence of paths to normalize

163

164 Returns:

165 List of normalized paths

166

167 Raises:

168 TypeError: If searchpath is not a valid type

169 ValueError: If searchpath is empty or contains invalid paths

170 """

171 try:

172 # Handle single values by converting to list

173 if isinstance(searchpath, str | AsyncPath):

174 searchpath_list = [searchpath]

175 else:

176 searchpath_list = list(searchpath)

177 if not searchpath_list:

178 raise ValueError("searchpath cannot be empty")

179

180 # Validate each path in the sequence

181 self._validate_sequence_paths(searchpath_list)

182 return searchpath_list

183 except (TypeError, ValueError) as e:

184 if isinstance(e, ValueError):

185 raise

186 raise TypeError(

187 "searchpath must be a string, AsyncPath, or sequence of strings/AsyncPaths"

188 ) from e

189

190 def _validate_sequence_paths(self, searchpath_list: list[AsyncPath | str]) -> None:

191 """Validate each path in a sequence.

192

193 Args:

194 searchpath_list: List of paths to validate

195

196 Raises:

197 TypeError: If any path is not a valid type

198 """

199 for i, path in enumerate(searchpath_list):

200 if not (isinstance(path, str) or hasattr(path, "parts")):

201 raise TypeError(

202 f"searchpath item {i} must be a string or AsyncPath, got {type(path)}"

203 )

204

205 def _convert_to_async_paths(

206 self, searchpath_list: list[AsyncPath | str]

207 ) -> list[AsyncPath]:

208 """Convert a list of paths to AsyncPath objects.

209

210 Args:

211 searchpath_list: List of paths to convert

212

213 Returns:

214 List of AsyncPath objects

215

216 Raises:

217 ValueError: If any path is an empty string

218 """

219 async_paths: list[AsyncPath] = []

220 for path in searchpath_list:

221 if isinstance(path, str):

222 if not path.strip():

223 raise ValueError("Empty string paths are not allowed")

224 async_paths.append(AsyncPath(path))

225 elif hasattr(path, "parts"):

226 # Already an AsyncPath-like object, convert to ensure type safety

227 async_paths.append(AsyncPath(str(path)))

228 else:

229 # Fallback: convert to string then AsyncPath

230 async_paths.append(AsyncPath(str(path)))

231 return async_paths

232

233 def _ensure_initialized(self) -> None:

234 """Ensure the loader is properly initialized.

235

236 This method provides thread-safe lazy initialization. Subclasses

237 should override _perform_initialization() instead of this method.

238

239 Thread Safety:

240 Uses a reentrant lock (RLock) to ensure thread-safe initialization

241 even when called from multiple threads concurrently.

242

243 Raises:

244 RuntimeError: If initialization fails

245 """

246 if not self._initialized:

247 # Thread-safe initialization check

248 import threading

249

250 if not hasattr(self, "_init_lock") or self._init_lock is None:

251 self._init_lock = threading.RLock()

252

253 with self._init_lock:

254 if not self._initialized:

255 try:

256 self._perform_initialization()

257 self._initialized = True

258 except Exception:

259 # Ensure we don't leave loader in partially initialized state

260 self._initialized = False

261 # Re-raise the original exception to maintain test compatibility

262 raise

263

264 def _perform_initialization(self) -> None:

265 """Perform the actual initialization work.

266

267 This method should be overridden by subclasses that need

268 custom initialization logic. It is called within a thread-safe

269 context by _ensure_initialized().

270

271 Subclasses should implement their initialization logic here

272 rather than overriding _ensure_initialized() directly.

273

274 Raises:

275 Exception: Any initialization-specific exceptions

276 """

277 pass

278

279 @abstractmethod

280 async def get_source_async(

281 self, environment: "AsyncEnvironment", name: str

282 ) -> SourceType:

283 """Get template source asynchronously.

284

285 Must be implemented by subclasses.

286

287 Args:

288 environment: The async environment instance

289 name: Template name to load

290

291 Returns:

292 Tuple of (source, filename, uptodate_func)

293

294 Raises:

295 TemplateNotFound: If template cannot be found

296 """

297 raise NotImplementedError("Subclasses must implement get_source_async")

298

299 async def list_templates_async(self) -> list[str]:

300 """List all available templates asynchronously.

301

302 Default implementation raises TypeError. Override in subclasses

303 that support template listing.

304

305 Returns:

306 List of template names

307

308 Raises:

309 TypeError: If listing is not supported by this loader

310 """

311 raise TypeError("this loader cannot iterate over all templates")

312

313 @internalcode

314 async def load_async(

315 self,

316 environment: "AsyncEnvironment",

317 name: str,

318 env_globals: dict[str, t.Any] | None = None,

319 ) -> "Template":

320 """Load template asynchronously using get_source_async.

321

322 This method orchestrates the complete template loading process:

323 1. Get template source via get_source_async()

324 2. Handle bytecode caching if available

325 3. Compile template if not cached

326 4. Create and return Template instance

327

328 Args:

329 environment: The async environment instance

330 name: Template name to load

331 env_globals: Global variables for the template (optional)

332

333 Returns:

334 Compiled Template object ready for rendering

335

336 Raises:

337 TemplateNotFound: If template cannot be found

338 RuntimeError: If compilation or loading fails

339 """

340 if env_globals is None:

341 env_globals = {}

342

343 # Import TemplateNotFound here to avoid circular imports

344

345 # Validate inputs and get template source

346 template_data = await self._prepare_template_loading_data(environment, name)

347

348 # Handle bytecode cache and compilation

349 code = await self._handle_template_compilation(environment, template_data)

350

351 # Create template instance

352 template: Template = environment.template_class.from_code(

353 environment,

354 code,

355 env_globals,

356 template_data.uptodate,

357 )

358

359 return template

360

361 async def _prepare_template_loading_data(

362 self, environment: "AsyncEnvironment", name: str

363 ) -> TemplateDataType:

364 """Prepare template loading data including validation and source retrieval.

365

366 Args:

367 environment: The async environment instance

368 name: Template name to load

369

370 Returns:

371 Named tuple containing source, path, and uptodate function

372

373 Raises:

374 TemplateNotFound: If template cannot be found or is invalid

375 """

376 from jinja2.exceptions import TemplateNotFound

377

378 # Validate inputs

379 if not name or not name.strip():

380 raise TemplateNotFound("Template name cannot be empty")

381

382 # Get template source

383 try:

384 source, path, uptodate = await self.get_source_async(environment, name)

385 except Exception as e:

386 if isinstance(e, TemplateNotFound):

387 raise

388 raise TemplateNotFound(f"Failed to get template source: {e}") from e

389

390 # Normalize source to string

391 try:

392 source_str = source.decode("utf-8") if isinstance(source, bytes) else source

393 except UnicodeDecodeError as e:

394 raise TemplateNotFound(

395 f"Template {name} contains invalid UTF-8 encoding: {e}"

396 ) from e

397

398 return TemplateDataType(source, path, uptodate, source_str, name)

399

400 async def _handle_template_compilation(

401 self, environment: "AsyncEnvironment", template_data: TemplateDataType

402 ) -> t.Any:

403 """Handle template compilation with bytecode caching.

404

405 Args:

406 environment: The async environment instance

407 template_data: Template data including source and path

408

409 Returns:

410 Compiled code object

411

412 Raises:

413 Exception: If compilation fails

414 """

415 # Handle bytecode cache if available

416 bcc = environment.bytecode_cache

417 if bcc is not None:

418 return await self._handle_bytecode_cache(environment, template_data, bcc)

419 # No cache, compile directly

420 return environment.compile(

421 template_data.source_str, template_data.name, template_data.path

422 )

423

424 async def _handle_bytecode_cache(

425 self,

426 environment: "AsyncEnvironment",

427 template_data: TemplateDataType,

428 bcc: t.Any,

429 ) -> t.Any:

430 """Handle bytecode caching.

431

432 Args:

433 environment: The async environment instance

434 template_data: Template data including source and path

435 bcc: Bytecode cache

436

437 Returns:

438 Compiled code object

439

440 Raises:

441 Exception: If cache operations fail

442 """

443 try:

444 bucket = bcc.get_bucket(environment, template_data.name, template_data.path)

445

446 # Create checksum for bytecode caching

447 import hashlib

448

449 hashlib.sha256(template_data.source_str.encode("utf-8")).hexdigest()

450

451 # Try to get existing bytecode

452 code = bucket.code

453

454 if code is None:

455 # Compile template

456 code = environment.compile(

457 template_data.source_str, template_data.name, template_data.path

458 )

459 bucket.code = code

460 return code

461 except Exception:

462 # Fallback to direct compilation if cache fails

463 return environment.compile(

464 template_data.source_str, template_data.name, template_data.path

465 )

466

467 def _get_cache_manager(self, environment: "AsyncEnvironment") -> t.Any:

468 """Get the cache manager from the environment.

469

470 Args:

471 environment: The async environment instance

472

473 Returns:

474 The cache manager instance

475 """

476 return getattr(environment, "cache_manager", None)

477

478 @internalcode

479 def _handle_template_not_found(self, name: str) -> None:

480 """Helper method to raise TemplateNotFound with consistent messaging.

481

482 Args:

483 name: Template name that was not found

484

485 Raises:

486 TemplateNotFound: Always raised with appropriate message

487 """

488 from jinja2.exceptions import TemplateNotFound

489

490 raise TemplateNotFound(name)