Get messages by id, return sorted lists

This commit is contained in:
Joscha 2016-05-09 22:18:12 +02:00
parent 4933caa6a9
commit 6e845a50a3

View file

@ -1,3 +1,5 @@
import operator
from . import message
class Messages():
@ -6,8 +8,17 @@ class Messages():
"""
def __init__(self):
self.by_id = {}
self.by_parent = {}
self._by_id = {}
self._by_parent = {}
def _sort(self, msgs):
"""
_sort(messages) -> None
Sorts a list of messages by their id, in place.
"""
msgs.sort(key=operator.attrgetter("id"))
def add_from_data(self, data):
"""
@ -29,25 +40,28 @@ class Messages():
self.remove(mes.id)
self.by_id[mes.id] = mes
self._by_id[mes.id] = mes
if mes.parent:
if not mes.parent in self.by_parent:
self.by_parent[mes.parent] = []
self.by_parent[mes.parent].append(mes)
if not mes.parent in self._by_parent:
self._by_parent[mes.parent] = []
self._by_parent[mes.parent].append(mes)
def remove(self, mes):
def remove(self, mid):
"""
remove(message) -> None
remove(message_id) -> None
Remove a message from the structure.
"""
if mes.id in self.by_id:
self.by_id.pop(mes.id)
if mes.parent and mes in self.get_children(mes.parent):
self.by_parent[mes.parent].remove(mes)
mes = self.get(mid)
if mes:
if mes.id in self._by_id:
self._by_id.pop(mes.id)
parent = self.get_parent(mes.id)
if parent and mes in self.get_children(parent.id):
self._by_parent[mes.parent].remove(mes)
def get(self, mid):
"""
@ -56,24 +70,39 @@ class Messages():
Returns the message with the given id, if found.
"""
if mid in self.by_id:
return self.by_id[mid]
if mid in self._by_id:
return self._by_id[mid]
def get_parent(self, mes):
def get_parent(self, mid):
"""
get_parent(message) -> str
get_parent(message_id) -> str
Returns the message's parent.
Returns None if no parent was found.
Returns the message's parent, if found.
"""
return self.get(mes.parent)
mes = self.get(mid)
if mes:
return self.get(mes.parent)
def get_children(self, mes):
def get_children(self, mid):
"""
get_children(message) -> list
get_children(message_id) -> list
Returns a list of children of the given message.
Returns a sorted list of children of the given message, if found.
"""
return self.by_parent[mes.id]
if mid in self._by_parent:
children = self._by_parent[mid][:]
self._sort(children)
return children
def get_top_level(self):
"""
get_top_level() -> list
Returns a sorted list of top-level messages.
"""
msgs = [self.get(mid) for mid in self._by_id if not self.get(mid).parent]
self._sort(msgs)
return msgs