|
28 | 28 | import java.util.ArrayList;
|
29 | 29 | import java.util.HashSet;
|
30 | 30 | import java.util.List;
|
| 31 | +import java.util.Stack; |
31 | 32 | import java.util.function.Function;
|
32 | 33 |
|
33 | 34 | import org.antlr.v4.runtime.ParserRuleContext;
|
@@ -55,7 +56,7 @@ public final class ScopeTranslator<T> extends Python3BaseVisitor<T> {
|
55 | 56 | private final ArrayList<String> possibleCellIdentifiers = new ArrayList<>();
|
56 | 57 | private final ArrayList<ScopeInfo> possibleCellScopes = new ArrayList<>();
|
57 | 58 |
|
58 |
| - private ScopeInfo currentGeneratorScope = null; |
| 59 | + private Stack<ScopeInfo> currentGeneratorScope = new Stack<>(); |
59 | 60 |
|
60 | 61 | public ScopeTranslator(ParserErrorCallback errors, TranslationEnvironment environment, boolean interactive, FrameDescriptor curInlineLocals) {
|
61 | 62 | this.errors = errors;
|
@@ -354,32 +355,37 @@ public T visitComp_for(Python3Parser.Comp_forContext ctx) {
|
354 | 355 | public T visitOr_test(Python3Parser.Or_testContext ctx) {
|
355 | 356 | boolean pushedCurrentGeneratorScope = false;
|
356 | 357 | if (ctx.getParent() instanceof Python3Parser.Comp_forContext) {
|
357 |
| - if (currentGeneratorScope == null && environment.getCurrentScopeLoopCount() == 1) { |
| 358 | + if (currentGeneratorScope.peek() == null && environment.getCurrentScopeLoopCount() == 1) { |
358 | 359 | // the generator iterator needs to be early evaluated in the parent scope
|
359 |
| - currentGeneratorScope = environment.pushCurentScope(); |
| 360 | + currentGeneratorScope.pop(); |
| 361 | + currentGeneratorScope.push(environment.pushCurentScope()); |
360 | 362 | pushedCurrentGeneratorScope = true;
|
361 | 363 | }
|
362 | 364 | }
|
363 | 365 | try {
|
364 | 366 | return super.visitOr_test(ctx);
|
365 | 367 | } finally {
|
366 | 368 | if (ctx.getParent() instanceof Python3Parser.Comp_forContext) {
|
367 |
| - if (pushedCurrentGeneratorScope && currentGeneratorScope.getLoopCount() == 1) { |
| 369 | + if (pushedCurrentGeneratorScope) { |
| 370 | + ScopeInfo scopeInfo = currentGeneratorScope.pop(); |
368 | 371 | // restore the current scope
|
369 |
| - environment.popCurrentScope(currentGeneratorScope); |
370 |
| - currentGeneratorScope = null; |
| 372 | + environment.popCurrentScope(scopeInfo); |
| 373 | + currentGeneratorScope.push(null); |
371 | 374 | }
|
372 | 375 | }
|
373 | 376 | }
|
374 | 377 | }
|
375 | 378 |
|
376 | 379 | private T visitGenerator(ParserRuleContext ctx, Python3Parser.Comp_forContext compctx, Function<ParserRuleContext, T> block) {
|
377 | 380 | compctx.scope = environment.createScope(ctx, ScopeKind.Generator);
|
| 381 | + currentGeneratorScope.push(null); |
378 | 382 | try {
|
379 | 383 | return block.apply(ctx);
|
380 | 384 | } finally {
|
381 |
| - if (currentGeneratorScope == null) { |
| 385 | + if (currentGeneratorScope.pop() == null) { |
382 | 386 | environment.leaveScope();
|
| 387 | + } else { |
| 388 | + throw new IllegalStateException("why did the currentGeneratorScope leak?"); |
383 | 389 | }
|
384 | 390 | }
|
385 | 391 | }
|
|
0 commit comments