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 from . import message
class Messages(): class Messages():
@ -6,8 +8,17 @@ class Messages():
""" """
def __init__(self): def __init__(self):
self.by_id = {} self._by_id = {}
self.by_parent = {} 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): def add_from_data(self, data):
""" """
@ -29,25 +40,28 @@ class Messages():
self.remove(mes.id) self.remove(mes.id)
self.by_id[mes.id] = mes self._by_id[mes.id] = mes
if mes.parent: if mes.parent:
if not mes.parent in self.by_parent: if not mes.parent in self._by_parent:
self.by_parent[mes.parent] = [] self._by_parent[mes.parent] = []
self.by_parent[mes.parent].append(mes) 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. Remove a message from the structure.
""" """
if mes.id in self.by_id: mes = self.get(mid)
self.by_id.pop(mes.id) if mes:
if mes.id in self._by_id:
if mes.parent and mes in self.get_children(mes.parent): self._by_id.pop(mes.id)
self.by_parent[mes.parent].remove(mes)
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): def get(self, mid):
""" """
@ -56,24 +70,39 @@ class Messages():
Returns the message with the given id, if found. Returns the message with the given id, if found.
""" """
if mid in self.by_id: if mid in self._by_id:
return self.by_id[mid] 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 the message's parent, if found.
Returns None if no parent was 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