ytsm/YtManagerApp/management/folders.py

from YtManagerApp.models import SubscriptionFolder, Subscription
from typing import Callable, Union, Any, Optional
from django.contrib.auth.models import User
import logging
from django.db.models.functions import Lower


def traverse_tree(root_folder_id: Optional[int], user: User, visit_func: Callable[[Union[SubscriptionFolder, Subscription]], Any]):
    data_collected = []

    def collect(data):
        if data is not None:
            data_collected.append(data)

    # Visit root
    if root_folder_id is not None:
        root_folder = SubscriptionFolder.objects.get(id = root_folder_id)
        collect(visit_func(root_folder))

    queue = [root_folder_id]
    visited = []

    while len(queue) > 0:
        folder_id = queue.pop()

        if folder_id in visited:
            logging.error('Found folder tree cycle for folder id %d.', folder_id)
            continue
        visited.append(folder_id)

        for folder in SubscriptionFolder.objects.filter(parent_id=folder_id, user=user).order_by(Lower('name')):
            collect(visit_func(folder))
            queue.append(folder.id)

        for subscription in Subscription.objects.filter(parent_folder_id=folder_id, user=user).order_by(Lower('name')):
            collect(visit_func(subscription))

    return data_collected