Rephrase async contextmanager DBConnection to use asynccontextmanager (#15538)

* Rephrase async contextmanager DBConnection to use asynccontextmanager
https://docs.python.org/3/library/contextlib.html#contextlib.asynccontextmanager

* drop tests.util.db_connection from typecheck exclusions list
This commit is contained in:
Adam Kelly 2023-06-16 14:42:14 -07:00 committed by GitHub
parent bb430d060b
commit d196247456
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
2 changed files with 13 additions and 16 deletions

View File

@ -142,7 +142,6 @@ tests.pools.test_wallet_pool_store
tests.simulation.test_simulation
tests.tools.test_run_block
tests.util.benchmark_cost
tests.util.db_connection
tests.util.generator_tools_testing
tests.util.key_tool
tests.util.test_full_block_utils

View File

@ -1,23 +1,21 @@
from __future__ import annotations
import tempfile
from contextlib import asynccontextmanager
from pathlib import Path
from typing import AsyncIterator
from chia.util.db_wrapper import DBWrapper2
class DBConnection:
def __init__(self, db_version: int) -> None:
self.db_version = db_version
async def __aenter__(self) -> DBWrapper2:
self.db_path = Path(tempfile.NamedTemporaryFile().name)
if self.db_path.exists():
self.db_path.unlink()
self._db_wrapper = await DBWrapper2.create(database=self.db_path, reader_count=4, db_version=self.db_version)
return self._db_wrapper
async def __aexit__(self, exc_t, exc_v, exc_tb) -> None:
await self._db_wrapper.close()
self.db_path.unlink()
@asynccontextmanager
async def DBConnection(db_version: int) -> AsyncIterator[DBWrapper2]:
db_path = Path(tempfile.NamedTemporaryFile().name)
if db_path.exists():
db_path.unlink()
_db_wrapper = await DBWrapper2.create(database=db_path, reader_count=4, db_version=db_version)
try:
yield _db_wrapper
finally:
await _db_wrapper.close()
db_path.unlink()