Converting pymysql Comment objects to tree recursively

Question

I am trying to create a comment system as a part of a hobby project, but I cant figure out how to recursively sort the Comment object once they have been fetched from the database. I'm using a relational database which has the following datamodel:

class Comment(Base):
    __tablename__ = 'comments'
    id = Column(Integer, primary_key=True)
    comment = Column(String(), nullable=False)
    user_id = Column(Integer, ForeignKey('users.id'), nullable=False)
    post_id = Column(Integer, ForeignKey('posts.id'), nullable=False)
    parent_id = Column(Integer, ForeignKey('comments.id'), nullable=False)

Once I've got the data from the database I need to sort these objects in a tree. Example input could for example be:

comments = [
            <models.Comment object at 0x104d80358>,
            <models.Comment object at 0x104d803c8>,
            <models.Comment object at 0x104d80470>, 
            <models.Comment object at 0x104d80518>,
            <models.Comment object at 0x104d805c0>,
            <models.Comment object at 0x104d80668>
           ]

And expected result could be:

comment_dict = {1: {'comment':<Comment.object>, 'children':[]},
               {2: {'comment':<Comment.object>, 'children':[<Comment.object>, ...]},
               {3: {'comment':<Comment.object>, 'children':[]},
               {4: {'comment':<Comment.object>, 'children':[<Comment.object>, ...]} ...

Any comment object can have an infinite amount of children. Almost like the comment system used on reddit and other similar social media sites. For the rendering I'm using flask and Jinja and could possibly do something like this which I found in the documentation:

<ul class="sitemap">
{%- for item in sitemap recursive %}
    <li><a href="{{ item.href|e }}">{{ item.title }}</a>
    {%- if item.children -%}
        <ul class="submenu">{{ loop(item.children) }}</ul>
    {%- endif %}</li>
{%- endfor %}

How I sort the data before doing this is what I'm not figuring out.

Answer 1

The very simple approach is this:

def comments_to_dict(comments):
    result = {}
    for comment in comments:
        result[comment.id] = {
            'comment': comment,
            'children': []
        }
    for comment in comments:
        result[comment.parent_id]['children'].append(comment)
    return result

So first you fill root elements with children as empty and then in second pass you fill children. This can be further improved by doing only one pass over comments:

def comments_to_dict(comments):
    result = {}
    for comment in comments:
        if comment.id in result:
            result[comment.id]['comment'] = comment
        else:
            result[comment.id] = {
                'comment': comment,
                'children': []
            }

        if comment.parent_id in result:
            result[comment.parent_id]['children'].append(comment)
        else:
            result[comment.parent_id] = {
                'children': [comment]
            }
    return result

The solution here matches the expected output you've shown us.

---

If however you want to have a real tree, then try this

def comments_to_dict(comments):
    index = {}
    for comment in comments:
        index[comment.id] = {
            'comment': comment,
            'children': []
        }
    for obj in index.itervalues():
        pid = obj['comment'].parent_id
        index[pid]['children'].append(obj)
    return index