python-app-rocketbot/rocketbot/collections.py

"""
Subclasses of list, set, and dict with special behaviors.
"""

from abc import ABCMeta, abstractmethod
from typing import Callable, Generic, Optional, TypeVar

# Abstract collections

# Dictionary key
K = TypeVar('K')
# Collection value
V = TypeVar('V')
# Element age
A = TypeVar('A')
# Age delta
D = TypeVar('D')

class AbstractMutableList(list[V], Generic[V], metaclass=ABCMeta):
	"""
	Abstract list with hooks for custom logic before and after mutation
	operations.
	"""
	@abstractmethod
	def pre_mutate(self):
		"""
		Called just prior to any mutation operation.
		"""
		raise NotImplementedError('subclass must implement pre_mutate')
	@abstractmethod
	def post_mutate(self):
		"""
		Called just after any mutating operation before returning the result
		to the caller.
		"""
		raise NotImplementedError('subclass must implement post_mutate')
	# Mutating method hooks
	def __delitem__(self, *args):
		self.pre_mutate()
		ret_val = super().__delitem__(*args)
		self.post_mutate()
		return ret_val
	def __iadd__(self, *args):
		self.pre_mutate()
		ret_val = super().__iadd__(*args)
		self.post_mutate()
		return ret_val
	def __imul__(self, *args):
		self.pre_mutate()
		ret_val = super().__imul__(*args)
		self.post_mutate()
		return ret_val
	def __setitem__(self, *args):
		self.pre_mutate()
		ret_val = super().__setitem__(*args)
		self.post_mutate()
		return ret_val
	def append(self, *args):
		self.pre_mutate()
		ret_val = super().append(*args)
		self.post_mutate()
		return ret_val
	def clear(self, *args):
		self.pre_mutate()
		ret_val = super().clear(*args)
		self.post_mutate()
		return ret_val
	def extend(self, *args):
		self.pre_mutate()
		ret_val = super().extend(*args)
		self.post_mutate()
		return ret_val
	def insert(self, *args):
		self.pre_mutate()
		ret_val = super().insert(*args)
		self.post_mutate()
		return ret_val
	def pop(self, *args):
		self.pre_mutate()
		ret_val = super().pop(*args)
		self.post_mutate()
		return ret_val
	def remove(self, *args):
		self.pre_mutate()
		ret_val = super().remove(*args)
		self.post_mutate()
		return ret_val
	def reverse(self, *args):
		self.pre_mutate()
		ret_val = super().reverse(*args)
		self.post_mutate()
		return ret_val
	def sort(self, *args):
		self.pre_mutate()
		ret_val = super().sort(*args)
		self.post_mutate()
		return ret_val

class AbstractMutableSet(set[V], Generic[V], metaclass=ABCMeta):
	"""
	Abstract set with hooks for custom logic before and after mutation
	operations.
	"""
	@abstractmethod
	def pre_mutate(self):
		"""
		Called just prior to any mutation operation.
		"""
		raise NotImplementedError('subclass must implement pre_mutate')
	@abstractmethod
	def post_mutate(self):
		"""
		Called just after any mutating operation before returning the result
		to the caller.
		"""
		raise NotImplementedError('subclass must implement post_mutate')
	# Mutating method hooks
	def __iand__(self, *args):
		self.pre_mutate()
		ret_val = super().__iand__(*args)
		self.post_mutate()
		return ret_val
	def __ior__(self, *args):
		self.pre_mutate()
		ret_val = super().__ior__(*args)
		self.post_mutate()
		return ret_val
	def __isub__(self, *args):
		self.pre_mutate()
		ret_val = super().__isub__(*args)
		self.post_mutate()
		return ret_val
	def __ixor__(self, *args):
		self.pre_mutate()
		ret_val = super().__ixor__(*args)
		self.post_mutate()
		return ret_val
	def add(self, *args):
		self.pre_mutate()
		ret_val = super().add(*args)
		self.post_mutate()
		return ret_val
	def clear(self, *args):
		self.pre_mutate()
		ret_val = super().clear(*args)
		self.post_mutate()
		return ret_val
	def difference_update(self, *args):
		self.pre_mutate()
		ret_val = super().difference_update(*args)
		self.post_mutate()
		return ret_val
	def discard(self, *args):
		self.pre_mutate()
		ret_val = super().discard(*args)
		self.post_mutate()
		return ret_val
	def intersection_update(self, *args):
		self.pre_mutate()
		ret_val = super().intersection_update(*args)
		self.post_mutate()
		return ret_val
	def pop(self, *args):
		self.pre_mutate()
		ret_val = super().pop(*args)
		self.post_mutate()
		return ret_val
	def remove(self, *args):
		self.pre_mutate()
		ret_val = super().remove(*args)
		self.post_mutate()
		return ret_val
	def symmetric_difference_update(self, *args):
		self.pre_mutate()
		ret_val = super().symmetric_difference_update(*args)
		self.post_mutate()
		return ret_val
	def update(self, *args):
		self.pre_mutate()
		ret_val = super().update(*args)
		self.post_mutate()
		return ret_val

class AbstractMutableDict(dict[K, V], Generic[K, V], metaclass=ABCMeta):
	"""
	Abstract dict with hooks for custom logic before and after mutation
	operations.
	"""
	@abstractmethod
	def pre_mutate(self):
		"""
		Called just prior to any mutation operation.
		"""
		raise NotImplementedError('subclass must implement pre_mutate')
	@abstractmethod
	def post_mutate(self):
		"""
		Called just after any mutating operation before returning the result
		to the caller.
		"""
		raise NotImplementedError('subclass must implement post_mutate')
	# Mutating method hooks
	def __delitem__(self, *args):
		self.pre_mutate()
		ret_val = super().__delitem__(*args)
		self.post_mutate()
		return ret_val
	def __ior__(self, *args):
		self.pre_mutate()
		ret_val = super().__ior__(*args)
		self.post_mutate()
		return ret_val
	def __setitem__(self, *args):
		self.pre_mutate()
		ret_val = super().__setitem__(*args)
		self.post_mutate()
		return ret_val
	def clear(self, *args):
		self.pre_mutate()
		ret_val = super().clear(*args)
		self.post_mutate()
		return ret_val
	def pop(self, *args):
		self.pre_mutate()
		ret_val = super().pop(*args)
		self.post_mutate()
		return ret_val
	def popitem(self, *args):
		self.pre_mutate()
		ret_val = super().popitem(*args)
		self.post_mutate()
		return ret_val
	def update(self, *args):
		self.pre_mutate()
		ret_val = super().update(*args)
		self.post_mutate()
		return ret_val

# Collections with limited number of elements

class SizeBoundList(AbstractMutableList[V], Generic[V, A]):
	"""
	Subclass of `list` that enforces a maximum number of elements.

	If the number of elements following a mutating operation exceeds
	`self.max_element_count`, then each element will be tested for its "age,"
	and the "oldest" elements will be removed until the total size is back
	within the limit. "Age" is determined via a provided lambda function. It is
	a value with arbitrary numeric value that is used comparatively to find
	the element with the smallest value.

	The `element_age` lambda takes two arguments: the element index and the
	element value. It must return values that can be compared to one another
	(e.g. int, float, datetime).

	`self.element_age` and `self.max_element_count` can be modified at runtime,
	however elements will only be discarded following the next mutating
	operation. Call `self.purge_old_elements()` to force resizing.
	"""
	def __init__(self,
			max_element_count: int,
			element_age: Callable[[int, V], A],
			*args, **kwargs):
		super().__init__(*args, **kwargs)
		self.element_age = element_age
		self.max_element_count = max_element_count
		self.is_culling = False
		self.post_mutate()

	def purge_old_elements(self):
		"""
		Manually purges elements. Purging usually occurs automatically, but
		this can be called if `self.max_element_count` or `self.element_age`
		are modified at runtime.
		"""
		self.post_mutate()

	def pre_mutate(self):
		pass

	def post_mutate(self):
		if self.is_culling:
			return
		self.is_culling = True
		while len(self) > self.max_element_count:
			oldest_age = None
			oldest_index = -1
			for i, elem in enumerate(self):
				age = self.element_age(i, elem)
				if oldest_age is None or age < oldest_age:
					oldest_age = age
					oldest_index = i
			self.pop(oldest_index)
		self.is_culling = False

	def copy(self):
		return SizeBoundList(self.max_element_count, self.element_age, super())

class SizeBoundSet(AbstractMutableSet[V], Generic[V, A]):
	"""
	Subclass of `set` that enforces a maximum number of elements.

	If the number of elements following a mutating operation exceeds
	`self.max_element_count`, then each element will be tested for its "age,"
	and the "oldest" elements will be removed until the total size is back
	within the limit. "Age" is determined via a provided lambda function. It is
	a value with arbitrary numeric value that is used comparatively to find
	the element with the smallest value.

	The `element_age` lambda takes one argument: the element value. It must
	return values that can be compared to one another (e.g. int, float,
	datetime).

	`self.element_age` and `self.max_element_count` can be modified at runtime,
	however elements will only be discarded following the next mutating
	operation. Call `self.purge_old_elements()` to force resizing.
	"""
	def __init__(self,
			max_element_count: int,
			element_age: Callable[[int, V], A],
			*args, **kwargs):
		super().__init__(*args, **kwargs)
		self.element_age = element_age
		self.max_element_count = max_element_count
		self.is_culling = False
		self.post_mutate()

	def purge_old_elements(self):
		"""
		Manually purges elements. Purging usually occurs automatically, but
		this can be called if `self.max_element_count` or `self.element_age`
		are modified at runtime.
		"""
		self.post_mutate()

	def pre_mutate(self):
		pass
	def post_mutate(self):
		if self.is_culling:
			return
		self.is_culling = True
		while len(self) > self.max_element_count:
			oldest_age = None
			oldest_elem = None
			for elem in self:
				age = self.element_age(elem)
				if oldest_age is None or age < oldest_age:
					oldest_age = age
					oldest_elem = elem
			self.remove(oldest_elem)
		self.is_culling = False

	def copy(self):
		return SizeBoundSet(self.max_element_count, self.element_age, super())

class SizeBoundDict(AbstractMutableDict[K, V], Generic[K, V, A]):
	"""
	Subclass of `dict` that enforces a maximum number of elements.

	If the number of elements following a mutating operation exceeds
	`self.max_element_count`, then each element will be tested for its "age,"
	and the "oldest" elements will be removed until the total size is back
	within the limit. "Age" is determined via a provided lambda function. It is
	a value with arbitrary numeric value that is used comparatively to find
	the element with the smallest value.

	The `element_age` lambda takes two arguments: the key and the value of a
	dict pair. It must return values that can be compared to one another (e.g.
	int, float, datetime).

	`self.element_age` and `self.max_element_count` can be modified at runtime,
	however elements will only be discarded following the next mutating
	operation. Call `self.purge_old_elements()` to force resizing.
	"""
	def __init__(self,
			max_element_count: int,
			element_age: Callable[[K, V], A],
			*args, **kwargs):
		super().__init__(*args, **kwargs)
		self.element_age = element_age
		self.max_element_count = max_element_count
		self.is_culling = False
		self.post_mutate()

	def purge_old_elements(self):
		"""
		Manually purges elements. Purging usually occurs automatically, but
		this can be called if `self.max_element_count` or `self.element_age`
		are modified at runtime.
		"""
		self.post_mutate()

	def pre_mutate(self):
		pass
	def post_mutate(self):
		if self.is_culling:
			return
		self.is_culling = True
		while len(self) > self.max_element_count:
			oldest_age = None
			oldest_key = None
			for key, value in self.items():
				age = self.element_age(key, value)
				if oldest_age is None or age < oldest_age:
					oldest_age = age
					oldest_key = key
			del self[oldest_key]
		self.is_culling = False

	def copy(self):
		return SizeBoundDict(self.max_element_count, self.element_age, super())

# Collections with limited age of elements

class AgeBoundList(AbstractMutableList[V], Generic[V, A, D]):
	"""
	Subclass of `list` that enforces a maximum "age" of elements.

	After each mutating operation, the minimum and maximum "age" of the elements
	are determined. If the span between the newest and oldest exceeds
	`self.max_age` then the oldest elements will be purged until that is
	no longer the case. "Age" is determined via a provided lambda function. It
	is a value with arbitrary numeric value that is used comparatively.

	The `element_age` lambda takes two arguments: the element index and the
	element value. It must return values that can be compared to one another
	and be added and subtracted (e.g. int, float, datetime). If the lambda
	returns `datetime`s, the `max_age` should be a `timedelta`.

	`self.element_age` and `self.max_age` can be modified at runtime,
	however elements will only be discarded following the next mutating
	operation. Call `self.purge_old_elements()` to force resizing.
	"""
	def __init__(self, max_age: D, element_age: Callable[[int, V], A], *args, **kwargs):
		super().__init__(*args, **kwargs)
		self.max_age = max_age
		self.element_age = element_age
		self.is_culling = False

	def purge_old_elements(self):
		"""
		Manually purges elements. Purging usually occurs automatically, but
		this can be called if `self.max_age` or `self.element_age` are modified
		at runtime.
		"""
		self.post_mutate()

	def pre_mutate(self):
		pass
	def post_mutate(self):
		if self.is_culling or len(self) <= 1:
			return
		self.is_culling = True
		min_age: Optional[A] = None
		max_age: Optional[A] = None
		ages: dict[int, A] = {}
		for i, elem in enumerate(self):
			age: A = self.element_age(i, elem)
			ages[i] = age
			if min_age is None or age < min_age:
				min_age = age
			if max_age is None or age > max_age:
				max_age = age
		cutoff: A = max_age - self.max_age
		if min_age >= cutoff:
			self.is_culling = False
			return
		for i in reversed(range(len(self))):
			if ages[i] < cutoff:
				del self[i]
		self.is_culling = False

	def copy(self):
		return AgeBoundList(self.max_age, self.element_age, super())

class AgeBoundSet(AbstractMutableSet[V], Generic[V, A, D]):
	"""
	Subclass of `set` that enforces a maximum "age" of elements.

	After each mutating operation, the minimum and maximum "age" of the elements
	are determined. If the span between the newest and oldest exceeds
	`self.max_age` then the oldest elements will be purged until that is
	no longer the case. "Age" is determined via a provided lambda function. It
	is a value with arbitrary numeric value that is used comparatively.

	The `element_age` lambda takes one argument: the element value. It must
	return values that can be compared to one another and be added and
	subtracted (e.g. int, float, datetime). If the lambda returns `datetime`s,
	the `max_age` should be a `timedelta`.

	`self.element_age` and `self.max_age` can be modified at runtime,
	however elements will only be discarded following the next mutating
	operation. Call `self.purge_old_elements()` to force resizing.
	"""
	def __init__(self, max_age: D, element_age: Callable[[int, V], A], *args, **kwargs):
		super().__init__(*args, **kwargs)
		self.max_age = max_age
		self.element_age = element_age
		self.is_culling = False

	def purge_old_elements(self):
		"""
		Manually purges elements. Purging usually occurs automatically, but
		this can be called if `self.max_age` or `self.element_age` are modified
		at runtime.
		"""
		self.post_mutate()

	def pre_mutate(self):
		pass
	def post_mutate(self):
		if self.is_culling or len(self) <= 1:
			return
		self.is_culling = True
		min_age = None
		max_age = None
		ages = {}
		for elem in self:
			age = self.element_age(elem)
			ages[elem] = age
			if min_age is None or age < min_age:
				min_age = age
			if max_age is None or age > max_age:
				max_age = age
		cutoff = max_age - self.max_age
		if min_age >= cutoff:
			self.is_culling = False
			return
		for elem, age in ages.items():
			if age < cutoff:
				self.remove(elem)
		self.is_culling = False

	def copy(self):
		return AgeBoundSet(self.max_age, self.element_age, super())

class AgeBoundDict(AbstractMutableDict[K, V], Generic[K, V, A, D]):
	"""
	Subclass of `dict` that enforces a maximum "age" of elements.

	After each mutating operation, the minimum and maximum "age" of the elements
	are determined. If the span between the newest and oldest exceeds
	`self.max_age` then the oldest elements will be purged until that is
	no longer the case. "Age" is determined via a provided lambda function. It
	is a value with arbitrary numeric value that is used comparatively.

	The `element_age` lambda takes two arguments: the key and value of a pair.
	It must return values that can be compared to one another and be added and
	subtracted (e.g. int, float, datetime). If the lambda returns `datetime`s,
	the `max_age` should be a `timedelta`.

	`self.element_age` and `self.max_age` can be modified at runtime,
	however elements will only be discarded following the next mutating
	operation. Call `self.purge_old_elements()` to force resizing.
	"""
	def __init__(self, max_age: D, element_age: Callable[[int, V], A], *args, **kwargs):
		super().__init__(*args, **kwargs)
		self.max_age = max_age
		self.element_age = element_age
		self.is_culling = False

	def purge_old_elements(self):
		"""
		Manually purges elements. Purging usually occurs automatically, but
		this can be called if `self.max_age` or `self.element_age` are modified
		at runtime.
		"""
		self.post_mutate()

	def pre_mutate(self):
		pass
	def post_mutate(self):
		if self.is_culling or len(self) <= 1:
			return
		self.is_culling = True
		min_age = None
		max_age = None
		ages = {}
		for key, value in self.items():
			age = self.element_age(key, value)
			ages[key] = age
			if min_age is None or age < min_age:
				min_age = age
			if max_age is None or age > max_age:
				max_age = age
		cutoff = max_age - self.max_age
		if min_age >= cutoff:
			self.is_culling = False
			return
		for key, age in ages.items():
			if age < cutoff:
				del self[key]
		self.is_culling = False

	def copy(self):
		return AgeBoundDict(self.max_age, self.element_age, super())