bank.py

from __future__ import annotations

import uuid
import copy
from contextlib import contextmanager
from dataclasses import dataclass, field
from enum import IntEnum, auto

def restore(originals, frozens):
    for original, frozen in zip(originals, frozens):
        assert original.uuid == frozen.uuid
        original.amount = frozen.amount

@contextmanager
def atomic(*wallets, hook = None):
    frozen_wallets = [copy.deepcopy(wallet) for wallet in wallets]
    
    try:
        yield
    except AssertionError:
        message = "<internal server error>"
    except Exception as e:
        message = "<atomic operation cancelled>"
    else:
        message = None

    if message is not None:
        restore(wallets, frozen_wallets)
        if hook is not None:
            hook(message)
    

class TransactionStatus(IntEnum):
    IN_PROGRESS = auto()
    CANCELLED = auto()
    SUCCESS = auto()
    FAILED = auto()

@dataclass
class Transaction:
    amount: int
    source: Wallet
    destination: Wallet
    status: TransactionStatus = TransactionStatus.IN_PROGRESS
    uuid: uuid.UUID = field(default_factory=uuid.uuid4)

    def __post_init__(self):
        self.source.record(self)

    def cancel(self):
        self.status = TransactionStatus.CANCELLED
    
    def process(self):
        if self.status is not TransactionStatus.IN_PROGRESS:
            return None
        elif self.amount < 0:
            return self.fail("<negative transaction>")
        elif self.amount == 0:
            return self.fail("<optimized away>")
        elif self.source.amount - self.amount < 0:
            return self.fail("<insufficient funds>")

        with atomic(self.source, self.destination, hook=self.fail):
            self.source.amount -= self.amount
            self.destination.amount += self.amount
        
            assert self.source.amount >= 0
            assert self.destination.amount >= 0
        
        self.succeed()

    def succeed(self):
        self.destination.record(self)
        self.status = TransactionStatus.SUCCESS

    def fail(self, message):
        self.status = TransactionStatus.FAILED
    


@dataclass
class Wallet:
    uuid: uuid.UUID = field(default_factory=uuid.uuid4)
    amount: int = 0
    transactions: List[Transaction] = field(repr=False, default_factory=list)

    def record(self, transaction):
        self.transactions.append(transaction)

    @contextmanager
    def _make_transaction(self, destination, amount):
        transaction = Transaction(amount, self, destination)
        yield transaction
        transaction.process()

    def send(self, destination, amount):
        with self._make_transaction(amount, destination) as transaction:
            return transaction

@dataclass
class Actor:
    identifier: str
    wallet: Wallet = field(default_factory=Wallet)

    def send(self, to, amount):
        return self.wallet.send(amount, to.wallet)

system = Actor('system',
    wallet = Wallet(
        amount = 1000
    )
)

actor_1 = Actor('b')
actor_2 = Actor('x')

system.send(actor_1, 100)
print(actor_1.transactions)