Skip to content

Instantly share code, notes, and snippets.

@isidentical
Last active Aug 17, 2021
Embed
What would you like to do?
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)
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment