generateloader.py

from collections import defaultdict

from promise import Promise
from promise.dataloader import DataLoader


def generate_loader(Type, attr):
    """ Generate a Loader on the fly.
    """
    class Loader(DataLoader):
        def batch_load_fn(self, keys):
            results_by_ids = defaultdict(list)
            lookup = {"{0}__in".format(attr): keys}
            for result in Type._meta.model.objects.filter(**lookup).iterator():
                results_by_ids[getattr(result, attr)].append(result)
            return Promise.resolve([results_by_ids.get(id, []) for id in keys])

    return Loader
    
    
loader = generate_loader(CommentType, "article_id")()