__all__ evaluation

src/Analysis/Ast/Impl/Analyzer/Evaluation/ExpressionEval.Operators.cs

@@ -13,8 +13,11 @@
 // See the Apache Version 2.0 License for specific language governing
 // permissions and limitations under the License.
 
+using System;
+using System.Collections.Generic;
 using Microsoft.Python.Analysis.Modules;
 using Microsoft.Python.Analysis.Types;
+using Microsoft.Python.Analysis.Types.Collections;
 using Microsoft.Python.Analysis.Values;
 using Microsoft.Python.Parsing;
 using Microsoft.Python.Parsing.Ast;
@@ -123,6 +126,16 @@ private IMember GetValueFromBinaryOp(Expression expr) {
                 return left;
             }
 
+            if (binop.Operator == PythonOperator.Add
+                && left.GetPythonType()?.TypeId == BuiltinTypeId.List
+                && right.GetPythonType()?.TypeId == BuiltinTypeId.List) {
+
+                var leftVar = GetValueFromExpression(binop.Left) as IPythonCollection;
+                var rightVar = GetValueFromExpression(binop.Right) as IPythonCollection;
+
+                return PythonCollectionType.CreateConcatenatedList(Module.Interpreter, GetLoc(expr), leftVar?.Contents, rightVar?.Contents);
+            }
+
             return left.IsUnknown() ? right : left;
         }
     }

src/Analysis/Ast/Impl/Analyzer/ModuleWalker.cs

@@ -13,12 +13,15 @@
 // See the Apache Version 2.0 License for specific language governing
 // permissions and limitations under the License.
 
+using System;
+using System.Collections.Generic;
 using System.Diagnostics;
 using System.Linq;
 using Microsoft.Python.Analysis.Analyzer.Evaluation;
 using Microsoft.Python.Analysis.Documents;
 using Microsoft.Python.Analysis.Modules;
 using Microsoft.Python.Analysis.Types;
+using Microsoft.Python.Analysis.Types.Collections;
 using Microsoft.Python.Analysis.Values;
 using Microsoft.Python.Core;
 using Microsoft.Python.Core.Collections;
@@ -33,6 +36,7 @@ internal class ModuleWalker : AnalysisWalker {
 
         // A hack to use __all__ export in the most simple case.
         private int _allReferencesCount;
+        private bool _allIsUsable = true;
 
         public ModuleWalker(IServiceContainer services, IPythonModule module, PythonAst ast)
             : base(new ExpressionEval(services, module, ast)) {
@@ -47,6 +51,103 @@ public override bool Walk(NameExpression node) {
             return base.Walk(node);
         }
 
+        public override bool Walk(AugmentedAssignStatement node) {
+            HandleAugmentedAllAssign(node);
+            return base.Walk(node);
+        }
+
+        public override bool Walk(CallExpression node) {
+            HandleAllAppendExtend(node);
+            return base.Walk(node);
+        }
+
+        private void HandleAugmentedAllAssign(AugmentedAssignStatement node) {
+            if (!IsHandleableAll(node.Left)) {
+                return;
+            }
+
+            if (node.Right is ErrorExpression) {
+                return;
+            }
+
+            if (node.Operator != Parsing.PythonOperator.Add) {
+                _allIsUsable = false;
+                return;
+            }
+
+            var rightVar = Eval.GetValueFromExpression(node.Right);
+            var rightContents = (rightVar as IPythonCollection)?.Contents;
+
+            if (rightContents == null) {
+                _allIsUsable = false;
+                return;
+            }
+
+            ExtendAll(node.Left, rightContents);
+        }
+
+        private void HandleAllAppendExtend(CallExpression node) {
+            if (!(node.Target is MemberExpression me)) {
+                return;
+            }
+
+            if (!IsHandleableAll(me.Target)) {
+                return;
+            }
+
+            if (node.Args.Count == 0) {
+                return;
+            }
+
+            IReadOnlyList<IMember> contents = null;
+            var v = Eval.GetValueFromExpression(node.Args[0].Expression);
+            if (v == null) {
+                _allIsUsable = false;
+                return;
+            }
+
+            switch (me.Name) {
+                case "append":
+                    contents = new List<IMember>() { v };
+                    break;
+                case "extend":
+                    contents = (v as IPythonCollection)?.Contents;
+                    break;
+            }
+
+            if (contents == null) {
+                _allIsUsable = false;
+                return;
+            }
+
+            ExtendAll(node, contents);
+        }
+
+        private void ExtendAll(Node declNode, IReadOnlyList<IMember> values) {
+            Eval.LookupNameInScopes(AllVariableName, out var scope, LookupOptions.Normal);
+            if (scope == null) {
+                return;
+            }
+
+            var loc = Eval.GetLoc(declNode);
+
+            var allContents = (scope.Variables[AllVariableName].Value as IPythonCollection)?.Contents;
+
+            var list = PythonCollectionType.CreateConcatenatedList(Module.Interpreter, loc, allContents, values);
+            var source = list.IsGeneric() ? VariableSource.Generic : VariableSource.Declaration;
+
+            Eval.DeclareVariable(AllVariableName, list, source, loc);
+        }
+
+        private bool IsHandleableAll(Node node) {
+            // TODO: handle more complicated lvars
+            if (!(node is NameExpression ne)) {
+                return false;
+            }
+
+            return Eval.CurrentScope == Eval.GlobalScope && ne.Name == AllVariableName;
+        }
+
         public override bool Walk(PythonAst node) {
             Check.InvalidOperation(() => Ast == node, "walking wrong AST");
 
@@ -98,7 +199,7 @@ public void Complete() {
             SymbolTable.ReplacedByStubs.Clear();
             MergeStub();
 
-            if (_allReferencesCount == 1 && GlobalScope.Variables.TryGetVariable(AllVariableName, out var variable) && variable?.Value is IPythonCollection collection) {
+            if (_allIsUsable && _allReferencesCount >= 1 && GlobalScope.Variables.TryGetVariable(AllVariableName, out var variable) && variable?.Value is IPythonCollection collection) {
                 ExportedMemberNames = collection.Contents
                     .OfType<IPythonConstant>()
                     .Select(c => c.GetString())

src/Analysis/Ast/Impl/Types/Collections/PythonCollectionType.cs

@@ -18,6 +18,7 @@
 using System.Linq;
 using Microsoft.Python.Analysis.Values;
 using Microsoft.Python.Analysis.Values.Collections;
+using Microsoft.Python.Core;
 
 namespace Microsoft.Python.Analysis.Types.Collections {
     /// <summary>
@@ -99,6 +100,11 @@ public static IPythonCollection CreateList(IPythonInterpreter interpreter, Locat
             return new PythonCollection(collectionType, location, contents, flatten);
         }
 
+        public static IPythonCollection CreateConcatenatedList(IPythonInterpreter interpreter, LocationInfo location, params IReadOnlyList<IMember>[] manyContents) {
+            var contents = manyContents?.ExcludeDefault().SelectMany().ToList() ?? new List<IMember>();
+            return CreateList(interpreter, location, contents);
+        }
+
         public static IPythonCollection CreateTuple(IPythonInterpreter interpreter, LocationInfo location, IReadOnlyList<IMember> contents) {
             var collectionType = new PythonCollectionType(null, BuiltinTypeId.Tuple, interpreter, false);
             return new PythonCollection(collectionType, location, contents);

src/LanguageServer/Test/ImportsTests.cs

@@ -459,5 +459,174 @@ public async Task FromImport_ModuleAffectsPackage(string appCodeImport) {
             comps = cs.GetCompletions(analysis, new SourceLocation(2, 21));
             comps.Should().HaveLabels("X");
         }
+
+        [TestMethod, Priority(0)]
+        public async Task AllSimple() {
+            var module1Code = @"
+class A:
+    def foo(self):
+        pass
+    pass
+
+class B:
+    def bar(self):
+        pass
+    pass
+
+__all__ = ['A']
+";
+
+            var appCode = @"
+from module1 import *
+
+A().
+B().
+";
+
+            var module1Uri = TestData.GetTestSpecificUri("module1.py");
+            var appUri = TestData.GetTestSpecificUri("app.py");
+
+            var root = Path.GetDirectoryName(appUri.AbsolutePath);
+            await CreateServicesAsync(root, PythonVersions.LatestAvailable3X);
+            var rdt = Services.GetService<IRunningDocumentTable>();
+            var analyzer = Services.GetService<IPythonAnalyzer>();
+
+            rdt.OpenDocument(module1Uri, module1Code);
+
+            var app = rdt.OpenDocument(appUri, appCode);
+            await analyzer.WaitForCompleteAnalysisAsync();
+            var analysis = await app.GetAnalysisAsync(-1);
+
+            var cs = new CompletionSource(new PlainTextDocumentationSource(), ServerSettings.completion);
+            var comps = cs.GetCompletions(analysis, new SourceLocation(4, 5));
+            comps.Should().HaveLabels("foo");
+
+            comps = cs.GetCompletions(analysis, new SourceLocation(5, 5));
+            comps.Should().NotContainLabels("bar");
+        }
+
+        [DataRow(@"
+other = ['B']
+__all__ = ['A'] + other")]
+        [DataRow(@"
+other = ['B']
+__all__ = ['A']
+__all__ += other")]
+        [DataRow(@"
+other = ['B']
+__all__ = ['A']
+__all__.extend(other)")]
+        [DataRow(@"
+__all__ = ['A']
+__all__.append('B')")]
+        [DataTestMethod, Priority(0)]
+        public async Task AllComplex(string allCode) {
+            var module1Code = @"
+class A:
+    def foo(self):
+        pass
+    pass
+
+class B:
+    def bar(self):
+        pass
+    pass
+
+class C:
+    def baz(self):
+        pass
+    pass
+" + allCode;
+
+            var appCode = @"
+from module1 import *
+
+A().
+B().
+C().
+";
+
+            var module1Uri = TestData.GetTestSpecificUri("module1.py");
+            var appUri = TestData.GetTestSpecificUri("app.py");
+
+            var root = Path.GetDirectoryName(appUri.AbsolutePath);
+            await CreateServicesAsync(root, PythonVersions.LatestAvailable3X);
+            var rdt = Services.GetService<IRunningDocumentTable>();
+            var analyzer = Services.GetService<IPythonAnalyzer>();
+
+            rdt.OpenDocument(module1Uri, module1Code);
+
+            var app = rdt.OpenDocument(appUri, appCode);
+            await analyzer.WaitForCompleteAnalysisAsync();
+            var analysis = await app.GetAnalysisAsync(-1);
+
+            var cs = new CompletionSource(new PlainTextDocumentationSource(), ServerSettings.completion);
+            var comps = cs.GetCompletions(analysis, new SourceLocation(4, 5));
+            comps.Should().HaveLabels("foo");
+
+            comps = cs.GetCompletions(analysis, new SourceLocation(5, 5));
+            comps.Should().HaveLabels("bar");
+
+            comps = cs.GetCompletions(analysis, new SourceLocation(6, 5));
+            comps.Should().NotContainLabels("baz");
+        }
+
+        [DataRow(@"
+__all__ = ['A']
+__all__.something(A)")]
+        [DataRow(@"
+__all__ = ['A']
+__all__ *= ['B']")]
+        [DataRow(@"
+__all__ = ['A']
+__all__ += 1234")]
+        [DataRow(@"
+__all__ = ['A']
+__all__.extend(123)")]
+        [DataRow(@"
+__all__ = ['A']
+__all__.extend(nothing)")]
+        [DataTestMethod, Priority(0)]
+        public async Task AllUnsupported(string allCode) {
+            var module1Code = @"
+class A:
+    def foo(self):
+        pass
+    pass
+
+class B:
+    def bar(self):
+        pass
+    pass
+" + allCode;
+
+            var appCode = @"
+from module1 import *
+
+A().
+B().
+";
+
+            var module1Uri = TestData.GetTestSpecificUri("module1.py");
+            var appUri = TestData.GetTestSpecificUri("app.py");
+
+            var root = Path.GetDirectoryName(appUri.AbsolutePath);
+            await CreateServicesAsync(root, PythonVersions.LatestAvailable3X);
+            var rdt = Services.GetService<IRunningDocumentTable>();
+            var analyzer = Services.GetService<IPythonAnalyzer>();
+
+            rdt.OpenDocument(module1Uri, module1Code);
+
+            var app = rdt.OpenDocument(appUri, appCode);
+            await analyzer.WaitForCompleteAnalysisAsync();
+            var analysis = await app.GetAnalysisAsync(-1);
+
+            var cs = new CompletionSource(new PlainTextDocumentationSource(), ServerSettings.completion);
+            var comps = cs.GetCompletions(analysis, new SourceLocation(4, 5));
+            comps.Should().HaveLabels("foo");
+
+            comps = cs.GetCompletions(analysis, new SourceLocation(5, 5));
+            comps.Should().HaveLabels("bar");
+        }
     }
 }