from promise import Promise
from promise.dataloader import DataLoader
from .models import NewsTranslate, CategoryTranslate, Keyword

class NewsLoader(DataLoader):
    def batch_load_fn(self, keys):
        news = NewsTranslate.objects.select_related(
            'news', 
            'language',
            'category'
        ).prefetch_related(
            'keywords'
        ).filter(id__in=keys)
        return Promise.resolve([
            next((n for n in news if n.id == key), None)
            for key in keys
        ])

class CategoryLoader(DataLoader):
    def batch_load_fn(self, keys):
        categories = CategoryTranslate.objects.select_related(
            'category'
        ).filter(id__in=keys)
        return Promise.resolve([
            next((c for c in categories if c.id == key), None) 
            for key in keys
        ])

class KeywordLoader(DataLoader):
    def batch_load_fn(self, keys):
        keywords = Keyword.objects.filter(id__in=keys)
        return Promise.resolve([
            next((k for k in keywords if k.id == key), None)
            for key in keys
        ])