The problem
You want to create a context manager from a generator function.
The solution
import functools
from contextlib import contextmanager
from typing import Callable, Generator
def as_contextmanager(func: Callable) -> Callable:
"""Decorator to convert a generator function into a context manager."""
@functools.wraps(func)
def helper(*args, **kwargs):
return _GeneratorContextManager(func, args, kwargs)
return helper
class _GeneratorContextManager:
def __init__(self, func, args, kwargs):
self.gen = func(*args, **kwargs)
def __enter__(self):
try:
return next(self.gen)
except StopIteration:
raise RuntimeError("generator didn't yield")
def __exit__(self, exc_type, exc_val, exc_tb):
if exc_type is None:
try:
next(self.gen)
except StopIteration:
return
else:
raise RuntimeError("generator didn't stop")
else:
try:
self.gen.throw(exc_type, exc_val, exc_tb)
except (StopIteration, RuntimeError):
return TruePut it to work
@as_contextmanager
def database_connection(connection_string: str):
"""Manage database connection with automatic cleanup."""
import sqlite3
conn = sqlite3.connect(connection_string)
print("Database connected")
try:
yield conn
finally:
conn.close()
print("Database disconnected")
# Usage
with database_connection(':memory:') as conn:
cursor = conn.cursor()
cursor.execute("CREATE TABLE test (id INTEGER)")
print("Table created")