Типизированный декоратор с параметрами

Условие

Напиши типизированный декоратор с параметрами. Такой, что он добавляет аргумент к декорируемой функции. Пример:

python 

from threading import Lock

my_lock = Lock()

@locked(my_lock)
def function1(lock: Lock, arg: int) -> int:
    return arg + 1

@locked(my_lock)
def function2(lock: Lock, name: str, surname: str) -> str:
    return name + surname

# Внутри 'locked' выполняется примерно такой код:
# python
# with passed_lock:
#     return function(passed_lock, *args, **kwargs)

# Вопрос со звёздочкой: как из него сделать и декоратор, и контекстный менеджер?

# Вопрос со звёздочкой: как из него сделать и декоратор, и контекстный менеджер?
# Вопрос со звёздочкой: как из него сделать и декоратор, и контекстный менеджер?

def locked(): # TOOO

Ответ

ParamSpec сохраняет параметры вызова, Concatenate описывает скрытый первый аргумент lock. Класс с call, enter и exit совмещает оба интерфейса.

Вариант для Python 3.10+ использует LockType как тип обычной блокировки:

from _thread import LockType
from functools import wraps
from typing import Callable, Concatenate, ParamSpec, TypeVar
from types import TracebackType

P = ParamSpec('P')
R = TypeVar('R')

class locked:
    def __init__(self, lock: LockType):
        self.lock = lock

    def __call__(self, fn: Callable[Concatenate[LockType, P], R]
                 ) -> Callable[P, R]:
        @wraps(fn)
        def wrapper(*args: P.args, **kwargs: P.kwargs) -> R:
            with self.lock:
                return fn(self.lock, *args, **kwargs)
        return wrapper

    def __enter__(self) -> LockType:
        self.lock.acquire()
        return self.lock

    def __exit__(self, kind: type[BaseException] | None,
                 value: BaseException | None,
                 tb: TracebackType | None) -> None:
        self.lock.release()

Вызов function1(3) автоматически передаст lock внутрь исходной функции. Доступно и with locked(my_lock). Обычный Lock не реентерабелен: повторный захват той же блокировки внутри функции или вложенного контекста вызовет зависание.