Python Splitting a Generator Yield into Two Parts

user_2955541 0
Viewed 631

I have access to a generator that yields two values:

def get_document_values():
    docs = query_database()  # returns a cursor to database documents
    for doc in docs:
        # doc is a dictionary with ,say, {'x': 1, 'y': 99}
        yield doc['x'], doc['y']

I have another function, process_x, that I cannot change that can take a generator in as input that processes all of the x for all documents (if a tuple is yielded then it just processes the first element of the tuple and ignores the other elements):

X = process_x(get_document_values())  # This processes x but ignores y

However, I need to store all of the y values from the generator as well. My only solution is to execute get_document_values twice:

Y = [y for x,y in get_document_values()]  #Throw away x
X = process_x(get_document_values())      #Throw away y

This technically works but when there are many documents to process, it is possible that a new document will get inserted into the database and the lengths of X and Y will be different. There needs to be a one-to-one mapping between X and Y and I'd like to only have to call get_document_values once instead of twice.

I've considered something like:

Y = []

def process_y(doc_generator):
    global Y
    for x,y in doc_generator:
        Y.append(y)
        yield x

X = process_x(process_y(get_document_values()))

But:

  1. This doesn't feel pythonic
  2. Y needs to be declared as a global variable

I am hoping that there is a cleaner, more pythonic way to do this.

Update

In reality, get_document_values will return values of x that are too large to be collectively stored into memory and process_x is actually reducing that memory requirement. So, it is not possible to cache all of x. Caching all of y is fine though.

generator python-3.x python
4 Answers

You are caching all the values into a list already when calling:

all_values = [(x,y) for x,y in get_document_values()] #or list(get_document_values())

You can get an iterator to y values with:

Y = map(itemgetter(1), all_values)

And for x simple use:

X = process_x(map(itemgetter(0), all_values))

The other option is to separate the generator, for example:

def get_document_values(getter):
    docs = query_database()  # returns a cursor to database documents
    for doc in docs:
        # doc is a dictionary with ,say, {'x': 1, 'y': 99}
        yield getter(doc)

from operator import itemgetter
X = process_x(get_document_values(itemgetter("x")))
Y = list(get_document_values(itemgetter("y")))

This way you will have to do the query twice, if you find a way of do the query once and duplicate the cursor, then you can abstract it also:

def get_document_values(cursor, getter):
    for doc in cursor:
        # doc is a dictionary with ,say, {'x': 1, 'y': 99}
        yield getter(doc)
user_16951720

No need to save the data:

def process_entry(x, y):
    process_x((x,))
    return y

ys = itertools.starmap(process_entry, your_generator)

Just remember that only when you get a y value, its corresponding x value is processed.

If you beed both, return both as a tuple:

def process_entry(x, y):
    return next(process_x((x,))), y
user_9485500

You may want to use itertools.tee to make two iterators from one, then use one iterator for process_x and second for the other purpose

user_4500290

Probably not pythonic, but you can cheat a little bit if it is permissible to change the main generator a little and make use of its function attribute:

from random import randrange
def get_vals():
        # mock creation of a x/y dict list
        docs =[{k: k+str(randrange(50)) for k in ('x','y')} for _ in range(10)]
        # create a function list attribute
        get_vals.y = []
        for doc in docs:
            # store the y value into the attribute
            get_vals.y.append(doc['y'])
            yield doc['x'], doc['y']  
            # if doc['y'] is purely for storage, you  might opt to not yield it at all.

Test it out:

# mock the consuming of generator for process_x            
for i in get_vals():
    print(i)    
# ('x13', 'y9'), ('x15', 'y40'), ('x41', 'y49')...

# access the ys stored in get_val function attribute after consumption
print(get_vals.y)
# ['y9', 'y40', 'y49', ...]

# instantiate the generator a second time a la process_x...
for i in get_vals():
    print(i)
# ('x18', 'y0'), ('x6', 'y35'), ('x24', 'y45')...

# access the cached y again
print(get_vals.y)
# ['y0', 'y35', 'y45', ...]
  1. This basically caches your y values as the generator outputs its x for each call.
  2. And it eliminates your global keyword
  3. And you can be sure the x/y mapping is accurate.

Some might consider this a hack, but I'd like to think of this as a feature since everything in Python is an object it allows you to get away with this...

user_91833440
Related