Search code examples
pythonalgorithmfunctiongeneratorenumerate

I have a generator in my Python function; how can I return a modified list?


A little background: I've been trying to code a "Sieve of Eratosthenes" algorithm. At the request of some fine (and very patient) programmers on the StackOverflow Python chatroom, I did some reading on the enumerate() function and found a way to incorporate it into my code (yes, I'm very much a novice here). So far, my (workable; returns expected response) code looks like this:

def SieveErat(n):
    numbers = [False]*2+[True]*(n-1)

    for index, prime_candidate in enumerate(numbers):
        if prime_candidate == True:
            yield index
            for x in xrange(index*index, n, index):
                numbers[x] = False

primes = []

for x in SieveErat(150000):
    primes.append(x)

print primes[10002]

Needless to say, the enumerate() function makes coding this much, much less awkward than whatever nested loops I had before. However, I fear that I'm not understanding something about enumerate(), because when I tried to shorten this code by including the appending into the function, I kept getting errors- namely,

File "SievErat.py", line 13
        return numbers
SyntaxError: 'return' with argument inside generator

I've also tried appending all the True elements inside the list numbers to a initialized list primes, but found no luck.

Any hints or advice would be very much welcomed.


Solution

  • This has nothing to do with enumerate you are trying to return something in a generator which before python 3.3 is illegal, and from 3.3+ it means something entirely different.

    I'd recommend you leave your function a generator if you may use it without needing a list back, and if you want list results then just call list() on the return value:

    primes = list(SieveErat(150000)) #this replaces loop with .append
    

    But to understand what went wrong, if your function still has the yield statement in it then it must return a generator object, if you don't want it to return a generator object then remove the yield statement all together:

    def SieveErat(n):
        numbers = [False]*2+[True]*(n-1)
    
        for index, prime_candidate in enumerate(numbers):
            if prime_candidate == True:
                #yield index #no yield statement if you want to return the numbers list
                for x in xrange(index*index, n, index):
                    numbers[x] = False
    
        return numbers #return at end
    

    However this would then return a list of True and False instead of the numbers that are True, you can instead hold a separate list with all the prime numbers and .append to it everytime you would yield something:

    def SieveErat(n):
        numbers = [False]*2+[True]*(n-1)
        result = [] #start with no results
        for index, prime_candidate in enumerate(numbers):
            if prime_candidate == True:
                results.append(index) #instead of yield index
                for x in xrange(index*index, n, index):
                    numbers[x] = False
        return results
    

    But this feels like a step backwards from a generator, personally I'd just keep what you have posted and cast the result to a list.