ytsm/YtManagerApp/management/folders.py

38 lines
1.2 KiB
Python

from YtManagerApp.models import SubscriptionFolder, Subscription
from typing import Callable, Union, Any, Optional
from django.contrib.auth.models import User
import logging
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('name'):
collect(visit_func(folder))
queue.append(folder.id)
for subscription in Subscription.objects.filter(parent_folder_id=folder_id, user=user).order_by('name'):
collect(visit_func(subscription))
return data_collected