python call generator function from other function-CodePudding

For the below code

# A simple generator function
def infinite_sequence():
    num = 1
    while True:
        yield num
        num  = 1

aaa = infinite_sequence()
bbb = infinite_sequence()
ccc = infinite_sequence()

print(next(aaa))
print(next(aaa))
print(next(bbb))
print(next(bbb))
print(next(ccc))
print(next(ccc))

the output is: 1 2 1 2 1 2

When trying to call the same generator functions from another function, the output is not as expected

def switchAction(input):
  action_item = {
    "A": next(aaa),
    "B": next(bbb),
    "C": next(ccc)
  }
  return action_item.get(input)
  

print(switchAction("A"))
print(switchAction("A"))
print(switchAction("B"))
print(switchAction("B"))
print(switchAction("C"))
print(switchAction("C"))

The output is: 1 2 3 4 5 6

Why the counter is continuing across generator functions in the case of calling from another function? How to achieve the same output as that of first case for the second case above.

CodePudding user response：

The problem is you call next on all values every time you call switchAction, since you define the dict over and over again. A solution to your problem can be as follows:

# A simple generator function
def infinite_sequence():
    num = 1
    while True:
        yield num
        num  = 1

aaa = infinite_sequence()
bbb = infinite_sequence()
ccc = infinite_sequence()

action_item = {
    "A": aaa,
    "B": bbb,
    "C": ccc
} # create the dict one time, and don't call next within the definition.

def switchAction(name): # also input is a dangerous var name so change it to name
    return next(action_item.get(name))


print(switchAction("A"))
print(switchAction("A"))
print(switchAction("B"))
print(switchAction("B"))
print(switchAction("C"))
print(switchAction("C"))

CodePudding user response：

Every time you call switchAction function, it is redefining action_item with the next values of generators. So you can fix that issue using callbacks.

def switchAction(input):
    action_item = {
        "A": lambda: next(aaa),
        "B": lambda: next(bbb),
        "C": lambda: next(ccc),
    }
    _next = action_item.get(input)
    if callable(_next):
        return _next()