You can already do this with the walrus operator:

# ---
def test():
    for i in range(10):
        yield (square := i * i)

        yield square + 1

# ---

This adds some parenthesis to your second alternative.