Drop file content to conserve memory (#1133)

* Remove local variables

* Drop file content to save memory

* Cache PEP hints

* Test fixes

* Add options to keep data in memory

Author: Mikhail Arkhipov

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

@@ -26,8 +26,15 @@ namespace Microsoft.Python.Analysis.Analyzer.Evaluation {
     /// and types in a chain of scopes during analysis.
     /// </summary>
     internal sealed partial class ExpressionEval {
+        private const string _pepHintKey = "PEP Hint";
+
         public IPythonType GetTypeFromPepHint(Node node) {
+            if (Ast.TryGetAttribute(node, _pepHintKey, out var typeStringObject) && typeStringObject is string typeString) {
+                return GetTypeFromString(typeString);
+            }
+
             var location = GetLocationInfo(node);
+
             var content = (Module as IDocument)?.Content;
             if (string.IsNullOrEmpty(content) || !location.EndLine.HasValue) {
                 return null;
@@ -79,7 +86,9 @@ public IPythonType GetTypeFromPepHint(Node node) {
             }
 
             // Type alone is not a valid syntax, so we need to simulate the annotation.
-            var typeString = content.Substring(hintStart, i - hintStart).Trim();
+            typeString = content.Substring(hintStart, i - hintStart).Trim();
+            Ast.SetAttribute(node, _pepHintKey, typeString);
+
             return GetTypeFromString(typeString);
         }
 

src/Analysis/Ast/Impl/Analyzer/Symbols/FunctionEvaluator.cs

@@ -79,6 +79,11 @@ public override void Evaluate() {
                 if (ctor || annotationType.IsUnknown() || Module.ModuleType == ModuleType.User) {
                     // Return type from the annotation is sufficient for libraries and stubs, no need to walk the body.
                     FunctionDefinition.Body?.Walk(this);
+                    // For libraries remove declared local function variables to free up some memory.
+                    var optionsProvider = Eval.Services.GetService<IAnalysisOptionsProvider>();
+                    if (Module.ModuleType != ModuleType.User && optionsProvider?.Options.KeepLibraryLocalVariables != true) {
+                        ((VariableCollection)Eval.CurrentScope.Variables).Clear();
+                    }
                 }
             }
             Result = _function;

src/Analysis/Ast/Impl/Definitions/AnalysisOptions.cs

@@ -16,5 +16,18 @@
 namespace Microsoft.Python.Analysis {
     public class AnalysisOptions {
         public bool LintingEnabled { get; set; }
+
+        /// <summary>
+        /// Keep in memory information on local variables declared in
+        /// functions in libraries. Provides ability to navigate to
+        /// symbols used in function bodies in packages and libraries.
+        /// </summary>
+        public bool KeepLibraryLocalVariables { get; set; }
+
+        /// <summary>
+        /// Keep in memory AST of library source code. May somewhat
+        /// improve performance when library code has to be re-analyzed.
+        /// </summary>
+        public bool KeepLibraryAst { get; set; }
     }
 }

src/Analysis/Ast/Impl/Modules/PythonModule.cs

@@ -438,6 +438,10 @@ public void NotifyAnalysisComplete(IDocumentAnalysis analysis) {
                 // as declare additional variables, etc.
                 OnAnalysisComplete();
                 ContentState = State.Analyzed;
+
+                if (ModuleType != ModuleType.User) {
+                    _buffer.Reset(_buffer.Version, string.Empty);
+                }
             }
 
             // Do not report issues with libraries or stubs

src/Analysis/Ast/Impl/Values/VariableCollection.cs

@@ -122,5 +122,11 @@ internal void RemoveVariable(string name) {
                 _variables.Remove(name);
             }
         }
+
+        internal void Clear() {
+            lock (_syncObj) {
+                _variables.Clear();
+            }
+        }
     }
 }

src/LanguageServer/Impl/Definitions/ServerSettings.cs

@@ -29,6 +29,22 @@ public class PythonAnalysisOptions {
             public string[] warnings { get; private set; } = Array.Empty<string>();
             public string[] information { get; private set; } = Array.Empty<string>();
             public string[] disabled { get; private set; } = Array.Empty<string>();
+
+            public class AnalysisMemoryOptions {
+                /// <summary>
+                /// Keep in memory information on local variables declared in
+                /// functions in libraries. Provides ability to navigate to
+                /// symbols used in function bodies in packages and libraries.
+                /// </summary>
+                public bool keepLibraryLocalVariables;
+
+                /// <summary>
+                /// Keep in memory AST of library source code. May somewhat
+                /// improve performance when library code has to be re-analyzed.
+                /// </summary>
+                public bool keepLibraryAst;
+            }
+            public AnalysisMemoryOptions memory;
         }
         public readonly PythonAnalysisOptions analysis = new PythonAnalysisOptions();
 

src/LanguageServer/Impl/Implementation/Server.cs

@@ -212,6 +212,11 @@ private bool HandleConfigurationChanges(ServerSettings newSettings) {
                 return true;
             }
 
+            if(newSettings.analysis?.memory?.keepLibraryAst != oldSettings.analysis?.memory?.keepLibraryAst ||
+               newSettings.analysis?.memory?.keepLibraryLocalVariables != oldSettings.analysis?.memory?.keepLibraryLocalVariables) {
+                return true;
+            }
+
             return false;
         }
 

src/LanguageServer/Impl/LanguageServer.Configuration.cs

@@ -73,6 +73,11 @@ private void HandleDiagnosticsChanges(JToken pythonSection, LanguageServerSettin
 
             var linting = pythonSection["linting"];
             HandleLintingOnOff(_services, GetSetting(linting, "enabled", true));
+
+            var memory = analysis["memory"];
+            var optionsProvider = _services.GetService<IAnalysisOptionsProvider>();
+            optionsProvider.Options.KeepLibraryLocalVariables = GetSetting(memory, "keepLibraryLocalVariables", false);
+            optionsProvider.Options.KeepLibraryAst = GetSetting(memory, "keepLibraryAst", false);
         }
 
         internal static void HandleLintingOnOff(IServiceContainer services, bool linterEnabled) {

src/LanguageServer/Test/GoToDefinitionTests.cs

@@ -332,6 +332,9 @@ def foo(self):
             var mainPath = TestData.GetTestSpecificUri("main.py");
             var doc = rdt.OpenDocument(mainPath, code);
 
+            var analyzer = Services.GetService<IPythonAnalyzer>();
+            await analyzer.WaitForCompleteAnalysisAsync();
+
             var analysis = await doc.GetAnalysisAsync(Timeout.Infinite);
             var ds = new DefinitionSource(Services);
 

src/Parsing/Impl/Ast/PythonAst.cs

@@ -96,15 +96,15 @@ public override async Task WalkAsync(PythonWalkerAsync walker, CancellationToken
 
         public PythonLanguageVersion LanguageVersion { get; }
 
-        internal bool TryGetAttribute(Node node, object key, out object value) {
+        public bool TryGetAttribute(Node node, object key, out object value) {
             if (_attributes.TryGetValue(node, out var nodeAttrs)) {
                 return nodeAttrs.TryGetValue(key, out value);
             }
             value = null;
             return false;
         }
 
-        internal void SetAttribute(Node node, object key, object value) {
+        public void SetAttribute(Node node, object key, object value) {
             if (!_attributes.TryGetValue(node, out var nodeAttrs)) {
                 nodeAttrs = _attributes[node] = new Dictionary<object, object>();
             }