heapq.merge with key=

[Chris Rebert]

On Thu, May 7, 2009 at 2:23 PM, Kevin D. Smith <Kevin.Smith at sas.com> wrote:
> I need the behavior of heapq.merge to merge a bunch of results from a
> database.  I was doing this with sorted(itertools.chain(...), key=...), but
> I would prefer to do this with generators.  My issue is that I need the key=
> argument to sort on the correct field in the database.  heapq.merge doesn't
> have this argument and I don't understand the code enough to know if it's
> possible to add it.  Is this enhancement possible without drastically
> changing the current code?

I think so. Completely untested code:

def key(ob):
    #code here

class Keyed(object):
    def __init__(self, obj):
        self.obj = obj
    def __cmp__(self, other):
        return cmp(key(self.obj), key(other.obj))

def keyify(gen):
    for item in gen:
        yield Keyed(item)

def stripify(gen):
    for keyed in gen:
        yield keyed.obj

merged = stripify(merge(keyify(A), keyify(B), keyify(C))) #A,B,C being
the iterables


Cheers,
Chris
-- 
http://blog.rebertia.com