Handling dependencies in module persistence

src/Analysis/Ast/Impl/Analyzer/Definitions/IAnalyzable.cs

@@ -13,11 +13,18 @@
 // See the Apache Version 2.0 License for specific language governing
 // permissions and limitations under the License.
 
+using Microsoft.Python.Analysis.Dependencies;
+
 namespace Microsoft.Python.Analysis.Analyzer {
     /// <summary>
     /// Represents document that can be analyzed asynchronously.
     /// </summary>
     internal interface IAnalyzable {
+        /// <summary>
+        /// Returns object that can calculate dependencies of this entry.
+        /// </summary>
+        IDependencyProvider DependencyProvider { get; }
+
         /// <summary>
         /// Notifies document that analysis is about to begin.
         /// </summary>

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

@@ -16,6 +16,7 @@
 using System.Collections.Generic;
 using System.Diagnostics;
 using System.Linq;
+using System.Threading;
 using Microsoft.Python.Analysis.Analyzer.Evaluation;
 using Microsoft.Python.Analysis.Documents;
 using Microsoft.Python.Analysis.Modules;
@@ -32,30 +33,36 @@ namespace Microsoft.Python.Analysis.Analyzer {
     internal class ModuleWalker : AnalysisWalker {
         private const string AllVariableName = "__all__";
         private readonly IDocumentAnalysis _stubAnalysis;
+        private readonly CancellationToken _cancellationToken;
 
         // 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)
+        public ModuleWalker(IServiceContainer services, IPythonModule module, PythonAst ast, CancellationToken cancellationToken)
             : base(new ExpressionEval(services, module, ast)) {
             _stubAnalysis = Module.Stub is IDocument doc ? doc.GetAnyAnalysis() : null;
+            _cancellationToken = cancellationToken;
         }
 
         public override bool Walk(NameExpression node) {
             if (Eval.CurrentScope == Eval.GlobalScope && node.Name == AllVariableName) {
                 _allReferencesCount++;
             }
+
+            _cancellationToken.ThrowIfCancellationRequested();
             return base.Walk(node);
         }
 
         public override bool Walk(AugmentedAssignStatement node) {
             HandleAugmentedAllAssign(node);
+            _cancellationToken.ThrowIfCancellationRequested();
             return base.Walk(node);
         }
 
         public override bool Walk(CallExpression node) {
             HandleAllAppendExtend(node);
+            _cancellationToken.ThrowIfCancellationRequested();
             return base.Walk(node);
         }
 
@@ -146,6 +153,7 @@ private bool IsHandleableAll(Node node) {
 
         public override bool Walk(PythonAst node) {
             Check.InvalidOperation(() => Ast == node, "walking wrong AST");
+            _cancellationToken.ThrowIfCancellationRequested();
 
             // Collect basic information about classes and functions in order
             // to correctly process forward references. Does not determine
@@ -181,16 +189,20 @@ public override bool Walk(PythonAst node) {
 
         // Classes and functions are walked by their respective evaluators
         public override bool Walk(ClassDefinition node) {
+            _cancellationToken.ThrowIfCancellationRequested();
             SymbolTable.Evaluate(node);
             return false;
         }
 
         public override bool Walk(FunctionDefinition node) {
+            _cancellationToken.ThrowIfCancellationRequested();
             SymbolTable.Evaluate(node);
             return false;
         }
 
         public void Complete() {
+            _cancellationToken.ThrowIfCancellationRequested();
+
             SymbolTable.EvaluateAll();
             SymbolTable.ReplacedByStubs.Clear();
             MergeStub();
@@ -220,11 +232,13 @@ public void Complete() {
         /// of the definitions. Stub may contains those so we need to merge it in.
         /// </remarks>
         private void MergeStub() {
-            if (Module.ModuleType == ModuleType.User) {
+            _cancellationToken.ThrowIfCancellationRequested();
+
+            if (Module.ModuleType == ModuleType.User || Module.ModuleType == ModuleType.Stub) {
                 return;
             }
             // No stub, no merge.
-            if (_stubAnalysis == null) {
+            if (_stubAnalysis.IsEmpty()) {
                 return;
             }
             // TODO: figure out why module is getting analyzed before stub.
@@ -248,6 +262,18 @@ private void MergeStub() {
                 if (stubType.DeclaringModule is TypingModule && stubType.Name == "Any") {
                     continue;
                 }
+
+                if (sourceVar?.Source == VariableSource.Import &&
+                   sourceVar.GetPythonType()?.DeclaringModule.Stub != null) {
+                    // Keep imported types as they are defined in the library. For example,
+                    // 'requests' imports NullHandler as 'from logging import NullHandler'.
+                    // But 'requests' also declares NullHandler in its stub (but not in the main code)
+                    // and that declaration does not have documentation or location. Therefore avoid
+                    // taking types that are stub-only when similar type is imported from another
+                    // module that also has a stub.
+                    continue;
+                }
+
                 TryReplaceMember(v, sourceType, stubType);
             }
 
@@ -257,7 +283,7 @@ private void MergeStub() {
         private void TryReplaceMember(IVariable v, IPythonType sourceType, IPythonType stubType) {
             // If type does not exist in module, but exists in stub, declare it unless it is an import.
             // If types are the classes, take class from the stub, then add missing members.
-            // Otherwise, replace type from one from the stub.
+            // Otherwise, replace type by one from the stub.
             switch (sourceType) {
                 case null:
                     // Nothing in sources, but there is type in the stub. Declare it.
@@ -379,28 +405,26 @@ private static void TransferDocumentationAndLocation(IPythonType s, IPythonType
             if (s != d && s is PythonType src && d is PythonType dst) {
                 // If type is a class, then doc can either come from class definition node of from __init__.
                 // If class has doc from the class definition, don't stomp on it.
-                var transferDoc = true;
                 if (src is PythonClassType srcClass && dst is PythonClassType dstClass) {
                     // Higher lever source wins
                     if (srcClass.DocumentationSource == PythonClassType.ClassDocumentationSource.Class ||
                        (srcClass.DocumentationSource == PythonClassType.ClassDocumentationSource.Init && dstClass.DocumentationSource == PythonClassType.ClassDocumentationSource.Base)) {
                         dstClass.SetDocumentation(srcClass.Documentation);
-                        transferDoc = false;
                     }
-                }
-
-                // Sometimes destination (stub type) already has documentation. This happens when stub type 
-                // is used to augment more than one type. For example, in threading module RLock stub class
-                // replaces both  RLock function and _RLock class making 'factory' function RLock to look 
-                // like a class constructor. Effectively a single  stub type is used for more than one type
-                // in the source and two source types may have different documentation. Thus transferring doc
-                // from one source type affects documentation of another type. It may be better to clone stub
-                // type and separate instances for separate source type, but for now we'll just avoid stomping
-                // on the existing documentation. 
-                if (transferDoc && string.IsNullOrEmpty(dst.Documentation)) {
-                    var srcDocumentation = src.Documentation;
-                    if (!string.IsNullOrEmpty(srcDocumentation)) {
-                        dst.SetDocumentation(srcDocumentation);
+                } else {
+                    // Sometimes destination (stub type) already has documentation. This happens when stub type 
+                    // is used to augment more than one type. For example, in threading module RLock stub class
+                    // replaces both  RLock function and _RLock class making 'factory' function RLock to look 
+                    // like a class constructor. Effectively a single  stub type is used for more than one type
+                    // in the source and two source types may have different documentation. Thus transferring doc
+                    // from one source type affects documentation of another type. It may be better to clone stub
+                    // type and separate instances for separate source type, but for now we'll just avoid stomping
+                    // on the existing documentation. 
+                    if (string.IsNullOrEmpty(dst.Documentation)) {
+                        var srcDocumentation = src.Documentation;
+                        if (!string.IsNullOrEmpty(srcDocumentation)) {
+                            dst.SetDocumentation(srcDocumentation);
+                        }
                     }
                 }
 

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

@@ -238,7 +238,7 @@ internal void RaiseAnalysisComplete(int moduleCount, double msElapsed)
         private void AnalyzeDocument(in AnalysisModuleKey key, in PythonAnalyzerEntry entry, in ImmutableArray<AnalysisModuleKey> dependencies) {
             _analysisCompleteEvent.Reset();
             ActivityTracker.StartTracking();
-            _log?.Log(TraceEventType.Verbose, $"Analysis of {entry.Module.Name}({entry.Module.ModuleType}) queued");
+            _log?.Log(TraceEventType.Verbose, $"Analysis of {entry.Module.Name} ({entry.Module.ModuleType}) queued. Dependencies: {string.Join(", ", dependencies.Select(d => d.IsTypeshed ? $"{d.Name} (stub)" : d.Name))}");
 
             var graphVersion = _dependencyResolver.ChangeValue(key, entry, entry.IsUserOrBuiltin || key.IsNonUserAsDocument, dependencies);
 

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

@@ -19,15 +19,13 @@
 using System.Linq;
 using System.Threading;
 using System.Threading.Tasks;
-using Microsoft.Python.Analysis.Core.DependencyResolution;
 using Microsoft.Python.Analysis.Modules;
 using Microsoft.Python.Analysis.Types;
-using Microsoft.Python.Core;
 using Microsoft.Python.Core.Collections;
 using Microsoft.Python.Parsing.Ast;
 
 namespace Microsoft.Python.Analysis.Analyzer {
-    [DebuggerDisplay("{_module.Name}({_module.ModuleType})")]
+    [DebuggerDisplay("{_module.Name} : {_module.ModuleType}")]
     internal sealed class PythonAnalyzerEntry {
         private readonly object _syncObj = new object();
         private TaskCompletionSource<IDocumentAnalysis> _analysisTcs;
@@ -248,20 +246,13 @@ public bool Invalidate(IPythonModule module, PythonAst ast, int bufferVersion, i
         }
 
         private HashSet<AnalysisModuleKey> FindDependencies(IPythonModule module, PythonAst ast, int bufferVersion) {
-            if (_bufferVersion > bufferVersion) {
-                return new HashSet<AnalysisModuleKey>();
-            }
-
-            var walker = new DependencyWalker(module);
-            ast.Walk(walker);
-            var dependencies = walker.Dependencies;
-            dependencies.Remove(new AnalysisModuleKey(module));
+            var dependencyProvider = (module as IAnalyzable)?.DependencyProvider;
+            var dependencies = _bufferVersion <= bufferVersion && module is IAnalyzable analyzable && analyzable.DependencyProvider != null
+                ? dependencyProvider.GetDependencies()
+                : new HashSet<AnalysisModuleKey>();
             return dependencies;
         }
 
-        private static bool Ignore(IModuleManagement moduleResolution, string fullName, string modulePath)
-            => moduleResolution.BuiltinModuleName.EqualsOrdinal(fullName) || moduleResolution.IsSpecializedModule(fullName, modulePath);
-
         private void UpdateAnalysisTcs(int analysisVersion) {
             _analysisVersion = analysisVersion;
             if (_analysisTcs.Task.Status == TaskStatus.RanToCompletion) {
@@ -273,69 +264,6 @@ private void UpdateAnalysisTcs(int analysisVersion) {
             }
         }
 
-        private class DependencyWalker : PythonWalker {
-            private readonly IPythonModule _module;
-            private readonly bool _isTypeshed;
-            private readonly IModuleManagement _moduleResolution;
-            private readonly PathResolverSnapshot _pathResolver;
-
-            public HashSet<AnalysisModuleKey> Dependencies { get; }
-
-            public DependencyWalker(IPythonModule module) {
-                _module = module;
-                _isTypeshed = module is StubPythonModule stub && stub.IsTypeshed;
-                _moduleResolution = module.Interpreter.ModuleResolution;
-                _pathResolver = _isTypeshed
-                    ? module.Interpreter.TypeshedResolution.CurrentPathResolver
-                    : _moduleResolution.CurrentPathResolver;
-
-                Dependencies = new HashSet<AnalysisModuleKey>();
-
-                if (module.Stub != null) {
-                    Dependencies.Add(new AnalysisModuleKey(module.Stub));
-                }
-            }
-
-            public override bool Walk(ImportStatement import) {
-                var forceAbsolute = import.ForceAbsolute;
-                foreach (var moduleName in import.Names) {
-                    var importNames = ImmutableArray<string>.Empty;
-                    foreach (var nameExpression in moduleName.Names) {
-                        importNames = importNames.Add(nameExpression.Name);
-                        var imports = _pathResolver.GetImportsFromAbsoluteName(_module.FilePath, importNames, forceAbsolute);
-                        HandleSearchResults(imports);
-                    }
-                }
 
-                return false;
-            }
-
-            public override bool Walk(FromImportStatement fromImport) {
-                var imports = _pathResolver.FindImports(_module.FilePath, fromImport);
-                HandleSearchResults(imports);
-                if (imports is IImportChildrenSource childrenSource) {
-                    foreach (var name in fromImport.Names) {
-                        if (childrenSource.TryGetChildImport(name.Name, out var childImport)) {
-                            HandleSearchResults(childImport);
-                        }
-                    }
-                }
-
-                return false;
-            }
-
-            private void HandleSearchResults(IImportSearchResult searchResult) {
-                switch (searchResult) {
-                    case ModuleImport moduleImport when !Ignore(_moduleResolution, moduleImport.FullName, moduleImport.ModulePath):
-                        Dependencies.Add(new AnalysisModuleKey(moduleImport.FullName, moduleImport.ModulePath, _isTypeshed));
-                        return;
-                    case PossibleModuleImport possibleModuleImport when !Ignore(_moduleResolution, possibleModuleImport.PrecedingModuleFullName, possibleModuleImport.PrecedingModulePath):
-                        Dependencies.Add(new AnalysisModuleKey(possibleModuleImport.PrecedingModuleFullName, possibleModuleImport.PrecedingModulePath, _isTypeshed));
-                        return;
-                    default:
-                        return;
-                }
-            }
-        }
     }
 }