Python – decorating a instance method with mypy (static type checker)

I’m new using mypy as a static type checker for my python projects and I’m having troubles trying to define a decorator for a instance method, where I want to access the instance properties. What I want to do is to decorate some class methods that require a method specific call before calling to some other methods. This is the target class:

class DataBase:

    _path: str
    _connection: sqlite3.Connection | None
    _cursor: sqlite3.Cursor | None

    def __init__(self, path: str) -> None:
        self._path = path
        self._connection = None
        self._cursor = None

    def connect(self) -> None:
        self._connection = sqlite3.connect(
            self._path, detect_types=sqlite3.PARSE_DECLTYPES | sqlite3.PARSE_COLNAMES
        )
        self._connection.execute("PRAGMA foreign_keys = ON;")
        self._cursor = self._connection.cursor()

    @require_connection
    def add_input_data(self, inputData: InputTable) -> None:
        ...
        ...
        ...

    @require_connection
    def close(self) -> None:
        self._connection.close()
        self._connection = None
        self._cursor = None

The decorator "require_connection" is something like this (trying to set mypy static types):

P = ParamSpec("P")
T = TypeVar("T")


def require_connection(func: Callable[P, T]) -> Callable[P, T]:
    @wraps(func)
    def wrapper(*args: P.args, **kwargs: P.kwargs) -> T:
        if self._connection is None:
            raise Exception(
                f"Cannot use the method: {func.__name__} before connecting to the Database"
            )
        return func(*args, **kwargs)

    return wrapper

How can I manage to use the "self._connection" in the wrapper?

Thank you in advance!

>Solution :

This is a design problem masquerading as a typing problem. Fix the design, and the typing problem goes away.

I assume you use your class something like

d = Database("...")
d.connect()
d.add_input_data(...)
d.close()

This suggests that either Database should be a context manager, or that its connect method returns a context manager.

with Database("...") as db:
    db.add_input_data(...)

or

db = Database("...")
with db.connect() as cm:
   cm.add_input_data(...)

The __enter__ method of the context manager is what will call sqlite3.connect and create the cursor, and the __exit__ method will close the connection. Rather than methods like add_input_data checking at runtime if they have an active connection to use, they will be on a class that cannot be successfully instantiated without successfully opening a connection. For example,

class DataBase:
    _path: str

    def __init__(self, path: str) -> None:
        self._path = path

    def connect(self):
        return DatabaseConnection(self)


class DatabaseConnection:
    def __init__(self, db):
        self.db = db

    def __enter__(self):
        self._connection = sqlite3.connect(self.db._path, ...)
        self._cursor = self._connection.cursor()

    def __exit__(self, *args):
        self._connection.close()

    def add_input_data(self, input_data: InputTable) -> None:
        ...


d = DataBase('...')

with d.connect() as conn:
    conn.add_input_data(...)