Django nested QuerySets

I have a Django data model like this (data fields omitted):

class Atom(Model):
    pass

class State(Model):
    atom = ForeignKey(Atom)

class Transition(Model):
    atom = ForeignKey(Atom)
    upstate = ForeignKey(State,related_name='uptrans')
    lostate = ForeignKey(State,related_name='lotrans')

When I query, the fields to be restricted can be in either model, so it is easiest to query on Transition.objects.filter(...) since all fields in the other models can be reached through the foreign keys. Let's call the resulting QuerySet t.

Now what I want in addition is a QuerySet a of the Atom model that corresponds to t, which can be done like a = t.values('atom').distinct(). So far so good.

However, I also want each of the entries in a to have one attribute/field that holds the QuerySet for the States of this Atom, still re开发者_StackOverflowflecting the criteria on the original selection t, through either one of the upstate or lostate ForeignKeys.

I have created my QuerySet on States up to now by looping over t, adding the values('upstate_id') and values('lostate_id') to a Python set() for throwing out duplicates, and then querying States with this list. But then I cannot achieve the nested structure of States within Atoms.

Any suggestions on how to do this are welcome, if possible with unevaluated QuerySets, since I pass them not into a template but a generator (yield statements), which is a nice way of streaming large amounts of data.

I think the following function does what I describe above, but I am not sure if the loop with further filtering of the original QuerySet by atom is the right approach.

def getAtomsWithStates(t):
    atom_ids = set( t.values_list('atom_id',flat=True) )
    atoms = Atoms.objects.filter(pk__in=atom_ids)
    for atom in atoms:
        upstate_ids = t.filter(atom=atom).values_list('upstate_id',flat=True)
        lostate_ids = t.filter(atom=atom).values_list('lostate_id',flat=True)
        all_ids = set( upstate_ids + lostate_ids )

        # attach the new QuerySet to the entry in the outer one:
        atom.States = State.objects.filter(pk__in=all_ids)

    return atoms

Now I can do the nested loop that I need like this:

someAtoms = getAtomsWithStates( Transition.objects.filter(...) )
for atom in someAtoms:
    for state in atom.States:
        print state.field

But, once again, there might be a smarter solution for this and I'd surely be interested.

It is very nice that you have the understanding of sets. However, using SQL's In will not duplicate your data. Let's consider that for a moment. If I say, "Give me an atom that is in this list: (1, 2, 3, 3, 3, 4)," the database will return atoms 1, 2, 3, and 4. For simplification, I would not ask Python to perform the set arithmetic since the database should be able to handle it just fine. There are times to use set, but your scenario doesn't seem like one of them.

An alternative for you:

states = State.objects.filter(
    Q(pk__in=t.values_list('upstate', flat=True)) |
    Q(pk__in=t.values_list('lostate', flat=True)
)

Even so, it seems like your model could use some changes, but I don't fully understand what you're trying to accomplish. Notice how that in my alternative, I don't do anything with atoms. I use the Q object to be able to perform an OR operation, but you might be able to add a flag to the State model to indicate if it is high or low. Or you could use an M2M relation with a through table. And why would your transition and your state both be associated with an atom? You could just eliminate atom from Transition and get the atom from State like so:

atoms = Atom.objects.filter(
    pk__in=State.objects.filter(
        Q(pk__in=t.values_list('upstate', flat=True)) |
        Q(pk__in=t.values_list('lostate', flat=True)
    ).values_list('atom', flat=True)
)