txn_test.py

import random
from functools import partial
import inspect
from typing import Optional
from types import MethodType
from contextlib import ContextDecorator

class Session:
    pass

class BuildingMotif:
    def __init__(self, db: str):
        self.db = db
        self._wrapping_context: Optional["Session"] = None

    def __getattribute__(self, name):
        attr = object.__getattribute__(self, name)
        #attr = super().__getattribute__(name)
        ctx = object.__getattribute__(self, "_wrapping_context")
        if ctx is None:
            return attr
        if hasattr(attr, '__call__') and 'ctx' in inspect.signature(attr).parameters:
            def wrap(*args, **kwargs):
                kwargs.update({'ctx': ctx})
                return attr(*args, **kwargs)
            return wrap
        else:
            return attr


    # library methods
    def foo(self, n: int, ctx: Optional["Session"] = None) -> int:
        if ctx is None:
            print(f"foo({n}) happens inside single transaction")
        else:
            print(f"foo({n}) happens inside an external transaction")
        return random.randint(0, n)

    def bar(self, ctx: Optional["Session"] = None) -> str:
        if ctx is None:
            print(f"bar() happens inside single transaction")
        else:
            print(f"bar() happens inside an external transaction")
        return "Hello there!"

    def baz(self, ctx: Optional["Session"] = None):
        if ctx is None:
            print(f"baz() happens inside single transaction")
        else:
            print(f"baz() happens inside an external transaction")

class GroupedTransaction(ContextDecorator):
    def __init__(self, bm: BuildingMotif):
        self.bm = bm

    def __enter__(self):
        print("Opening new transaction")
        self.session = "new session"
        self.bm._wrapping_context = self

    def __exit__(self, exc_type, exc, exc_tb):
        print("committing transaction")
        self.bm._wrapping_context = None


if __name__ == '__main__':
    bm = BuildingMotif("my_db")
    print(bm.foo(3))
    print(bm.bar())

    with GroupedTransaction(bm):
        bm.foo(3)
        bm.bar()