Better scoping checks

Author: MikhailArkhipov

src/Analysis/Engine/Impl/ModuleAnalysis.cs

@@ -246,39 +246,66 @@ private static VariablesResult GetVariablesFromCallExpression(AnalysisUnit unit,
             return null;
         }
 
+        private class VariableScopePair : Tuple<IAnalysisVariable, int> {
+            public VariableScopePair(IAnalysisVariable v, int scopeLevel) : base(v, scopeLevel) { }
+            public IAnalysisVariable Variable => Item1;
+            public int ScopeLevel => Item2;
+            public VariableType VariableType => Variable.Type;
+            public ILocationInfo Location => Variable.Location;
+        }
+
         private VariablesResult GetVariablesFromNameExpression(Expression expr, AnalysisUnit unit, IScope scope) {
             if (!(expr is NameExpression name)) {
                 return null;
             }
 
-            var variables = Enumerable.Empty<IAnalysisVariable>();
             if (!scope.EnumerateTowardsGlobal.Any()) {
-                variables = _unit.State.BuiltinModule.GetDefinitions(name.Name).SelectMany(ToVariables);
-                return new VariablesResult(variables, unit.Tree);
+                var v = _unit.State.BuiltinModule.GetDefinitions(name.Name).SelectMany(ToVariables);
+                return new VariablesResult(v, unit.Tree);
             }
 
+            var variables = Enumerable.Empty<VariableScopePair>();
+            VariableScopePair mainDefinition = null;
+            var scopeLevel = 0;
             foreach (var s in scope.EnumerateTowardsGlobal) {
-                var scopeVariables = GetVariablesInScope(name, s).Distinct();
+                var scopeVariables = GetVariablesInScope(name, s)
+                    .Distinct()
+                    .Select(v => new VariableScopePair(v, scopeLevel))
+                    .ToArray();
+
                 variables = variables.Union(scopeVariables);
-                var args = scopeVariables.Where(v => IsFunctionArgument(s, v));
-                if (args.Any()) {
+
+                // If we already have definition and it is a function argument, then stop
+                // since function argument wins over definitions in the outer scope
+                mainDefinition = scopeVariables.FirstOrDefault(v => v.VariableType == VariableType.Definition && IsFunctionArgument(s, v.Variable));
+                if (mainDefinition != null) {
                     break;
                 }
+                scopeLevel++;
             }
 
-            // Now take outermost definition and treat inner ones (such as reassignments) as references
-            var others = variables.Where(v => v.Type == VariableType.Reference || v.Type == VariableType.Value);
+            // Now take innermost definition and treat inner ones (such as reassignments) as references
             var definitions = variables
-                .Where(v => v.Type == VariableType.Definition)
+                .Where(v => v.VariableType == VariableType.Definition)
                 .OrderBy(v => v.Location.Span.Start)
+                .Reverse()
                 .ToArray();
 
-            if (definitions.Length > 0) {
-                var defsToRefs = definitions.Skip(1).Select(v => new AnalysisVariable(v.Variable, VariableType.Reference, v.Location));
-                variables = definitions.Take(1).Concat(others.Concat(defsToRefs));
+            mainDefinition = mainDefinition ?? definitions.FirstOrDefault();
+            if (mainDefinition != null) {
+                // Drop definitions in outer scopes and convert those in inner scopes to references.
+                // Scope levels are numbered in reverse (X == main definition level, x+1 == one up).
+                var defsToRefs = definitions
+                    .Where(d => d != mainDefinition && d.ScopeLevel <= mainDefinition.ScopeLevel)
+                    .Select(v => new VariableScopePair(new AnalysisVariable(v.Variable.Variable, VariableType.Reference, v.Location), v.ScopeLevel));
+
+                var others = variables
+                            .Where(v => (v.VariableType == VariableType.Reference || v.VariableType == VariableType.Value) &&
+                                         v.ScopeLevel <= mainDefinition.ScopeLevel);
+                variables = new[] { mainDefinition }.Concat(others.Concat(defsToRefs));
             }
 
-            return new VariablesResult(variables, unit.Tree);
+            return new VariablesResult(variables.Select(v => v.Variable), unit.Tree);
         }
 
         private VariablesResult GetVariablesFromMemberExpression(Expression expr, AnalysisUnit unit, ExpressionEvaluator eval) {

src/Analysis/Engine/Test/AnalysisTest.cs

@@ -3154,6 +3154,27 @@ assert isinstance(a, int)
             }
         }
 
+        [TestMethod, Priority(0)]
+        public async Task References_ClassLocalVariable() {
+            var uri = TestData.GetDefaultModuleUri();
+            using (var server = await CreateServerAsync(PythonVersions.LatestAvailable2X)) {
+                var text = @"
+a = 1
+class B:
+    a = 2
+";
+                await server.SendDidOpenTextDocument(uri, text);
+                var references = await server.SendFindReferences(uri, 3, 5);
+                references.Should().OnlyHaveReferences(
+                    (uri, (3, 4, 3, 5), ReferenceKind.Definition)
+                );
+                references = await server.SendFindReferences(uri, 1, 1);
+                references.Should().OnlyHaveReferences(
+                    (uri, (1, 0, 1, 1), ReferenceKind.Definition)
+                );
+            }
+        }
+
         [TestMethod, Priority(0)]
         public async Task ListDictArgReferences() {
             var text = @"

src/Analysis/Engine/Test/LanguageServerTests.cs

@@ -449,7 +449,10 @@ public async Task CompletionForOverrideArgs() {
                 var u = await AddModule(s, "class A:\n    def bar(arg=None): pass\n\nclass B(A):\n    def b");
 
                 await AssertNoCompletion(s, u, new SourceLocation(2, 9));
-                await AssertCompletion(s, u, new[] { "bar(arg=None):\r\n    return super().bar()" }, new[] { "bar(arg = None):\r\n    return super().bar()" }, new SourceLocation(5, 10));
+                await AssertCompletion(s, u, 
+                    new[] { "bar(arg=None):\r\n    return super(B, arg).bar()" }, 
+                    new[] { "bar(arg = None):\r\n    return super().bar()" }, 
+                    new SourceLocation(5, 10));
             }
         }
 
@@ -762,7 +765,7 @@ class D:
             await AssertReferences(s, mod1, new SourceLocation(10, 2), expected, unexpected);
             await AssertReferences(s, mod2, new SourceLocation(17, 6), expected, unexpected);
 
-            await AssertReferences(s, mod1, new SourceLocation(8, 5), unexpected, expected);
+            await AssertReferences(s, mod1, new SourceLocation(8, 5), unexpected, Enumerable.Empty<string>());
 
             // a
             expected = new[] {
@@ -800,14 +803,11 @@ class D:
 
             // C.real
             expected = new[] {
-                "Reference;(3, 7) - (3, 11)",
-                "Definition;(7, 5) - (7, 9)"
-            };
-            unexpected = new[] {
                 "Definition;(11, 1) - (11, 5)",
-                "Definition;(14, 5) - (14, 9)"
+                "Reference;(3, 7) - (3, 11)",
+                "Reference;(7, 5) - (7, 9)"
             };
-            await AssertReferences(s, mod1, new SourceLocation(7, 8), expected, unexpected);
+            await AssertReferences(s, mod1, new SourceLocation(7, 8), expected, Enumerable.Empty<string>());
 
             // D.real
             expected = new[] {
@@ -1228,7 +1228,11 @@ public static async Task AssertReferences(Server s, TextDocumentIdentifier docum
                 }
             }, CancellationToken.None);
 
-            refs.Select(r => $"{r._kind ?? ReferenceKind.Reference};{r.range}").Should().Contain(contains).And.NotContain(excludes);
+            if (excludes.Any()) {
+                refs.Select(r => $"{r._kind ?? ReferenceKind.Reference};{r.range}").Should().Contain(contains).And.NotContain(excludes);
+            } else {
+                refs.Select(r => $"{r._kind ?? ReferenceKind.Reference};{r.range}").Should().Contain(contains);
+            }
         }
 
         public Task<ILanguageServerExtension> CreateAsync(IPythonLanguageServer server, IReadOnlyDictionary<string, object> properties, CancellationToken cancellationToken) => throw new NotImplementedException();