from pprint import pformat
from typing import List
from typing import Optional
from typing import Tuple
from pytezos.michelson.types.base import MichelsonType
[docs]class MichelsonStack:
def __init__(self, items: Optional[List[MichelsonType]] = None) -> None:
self.items = items or []
self.protected = 0
[docs] @classmethod
def from_items(cls, items: List[MichelsonType]) -> 'MichelsonStack':
return cls(items)
[docs] def protect(self, count: int) -> None:
if len(self.items) < count:
raise Exception(f'got {len(self.items)} items on the stack, want to protect {count}')
self.protected += count
[docs] def restore(self, count: int) -> None:
if self.protected < count:
raise Exception(f'want to restore {count} items, but only {self.protected} are protected')
self.protected -= count
[docs] def push(self, item: MichelsonType):
self.items.insert(self.protected, item)
[docs] def peek(self) -> MichelsonType:
if not self.items:
raise Exception('stack is empty')
return self.items[self.protected]
[docs] def pop(self, count: int) -> List[MichelsonType]:
if len(self.items) - self.protected < count:
raise Exception(f'got {len(self.items) - self.protected} items on the stack, want to pop {count}')
return [self.items.pop(self.protected) for _ in range(count)]
[docs] def pop1(self) -> MichelsonType:
(a,) = self.pop(count=1)
return a
[docs] def pop2(self) -> Tuple[MichelsonType, MichelsonType]:
a, b = self.pop(count=2)
return a, b
[docs] def pop3(self) -> Tuple[MichelsonType, MichelsonType, MichelsonType]:
a, b, c = self.pop(count=3)
return a, b, c
[docs] def clear(self) -> None:
self.items.clear()
self.protected = 0
[docs] def dump(self, count: int) -> Optional[List[MichelsonType]]:
if not self.items:
return None
count = min(count, len(self.items))
return self.items[:count]
def __len__(self) -> int:
return len(self.items)
def __repr__(self) -> str:
return pformat(self.items)