Executar cascade query para resolver dados de todos os campos referencefield em um document.

cascade_query.py

'''
Usando o mongoengine, eu tenho alguns documentos no meu models.py que contém referencefield...é um recurso bem útil,
mas ao buscar os objetos desse documento, nesse campo eu tenho apenas o ObjectId, ou seja, o driver não faz uma busca no estilo
cascata. A ideia foi escrever um algoritmo simples (talvez não muito pythonico) para identificar reference fields e 
recuperar os dados.
'''
#Um exemplo simples para models.py
from app import db #db = MongoEngine()


class Cidades(db.Document):
    '''
    Representa o documento cidades...
    '''
    nome = db.StringField()
    estado = db.StringField()


class Pessoa(db.Document):
    '''
    Representa o documento pessoa
    '''
    nome = db.StringField()
    idade = db.StringField()
    cidade = db.ReferenceField(Cidades)

'''
Se eu executar uma query sobre o document Pessoa, o campo cidade terá apenas um DBRef com um ObjectID.
algo como:
'''

from meu_pacote import models
pessoa = models.Pessoa.objects().first()
#pessoa.cidade é apenas uma referencia

'''
A função abaixo recebe um document e retorna um dicionário com todos os referencefields, ou seja, DBRef's substituidos
pelos seus dados de fato. Um efeito cascata...
Existe ainda um detalhe: o campo referenciado pode ser uma lista de referencefields - o que também é tratado na função abaixo.
'''

def cascade_query(document):
    from mongoengine import ReferenceField, ListField
    from bson import DBRef
    doc_copy = document._data.copy()
    for key, field in document._fields.items():
        if isinstance(field, ReferenceField):
            if isinstance(document._data[key], DBRef):
                ref_res = field.document_type_obj.objects(id=document._data[key].id).first()
                doc_copy[key] = cascade_query(ref_res)
            else:
                doc_copy[key] = None
        elif isinstance(field, ListField):
            doc_copy[key] = [cascade_query(item) for item in
                             document if isinstance(item, ListField)]
    return doc_copy