Return appropriate derived class when method returns self

Python/Product/Analysis/Analyzer/DDG.cs

@@ -593,7 +593,7 @@ public override bool Walk(WithStatement node) {
                 var ctxMgr = _eval.Evaluate(item.ContextManager);
                 var enter = ctxMgr.GetMember(node, _unit, node.IsAsync ? "__aenter__" : "__enter__");
                 var exit = ctxMgr.GetMember(node, _unit, node.IsAsync ? "__aexit__" : "__exit__");
-                var ctxt = enter.Call(node, _unit, ExpressionEvaluator.EmptySets, ExpressionEvaluator.EmptyNames).Resolve(_unit);
+                var ctxt = enter.Call(node, _unit, new[] { ctxMgr }, ExpressionEvaluator.EmptyNames).Resolve(_unit);
                 var exitRes = exit.Call(node, _unit, ExpressionEvaluator.EmptySets, ExpressionEvaluator.EmptyNames).Resolve(_unit);
                 if (node.IsAsync) {
                     ctxt = ctxt.Await(node, _unit);

Python/Product/Analysis/Interpreter/Ast/AstPythonModule.cs

@@ -265,11 +265,6 @@ private static string TryGetDocFromModuleInitFile(string filePath) {
                             return line.Substring(quote.Length, line.Length - 2 * quote.Length).Trim();
                         }
                         var sb = new StringBuilder();
-                        // Handle '''Text right here
-                        line = line.Substring(quote.Length).Trim();
-                        if (!string.IsNullOrEmpty(line)) {
-                            sb.AppendLine(line);
-                        }
                         while (true) {
                             line = sr.ReadLine();
                             if (line == null || line.EndsWithOrdinal(quote)) {

Python/Product/Analysis/ProjectEntry.cs

@@ -42,7 +42,7 @@ namespace Microsoft.PythonTools.Analysis {
         Justification = "Unclear ownership makes it unlikely this object will be disposed correctly")]
     internal sealed class ProjectEntry : IPythonProjectEntry, IAggregateableProjectEntry, IDocument {
         private AnalysisUnit _unit;
-        private TaskCompletionSource<ModuleAnalysis> _analysisTcs;
+        private TaskCompletionSource<ModuleAnalysis> _analysisTcs = new TaskCompletionSource<ModuleAnalysis>();
         private readonly SortedDictionary<int, DocumentBuffer> _buffers;
         private readonly ConcurrentQueue<WeakReference<ReferenceDict>> _backReferences = new ConcurrentQueue<WeakReference<ReferenceDict>>();
         internal readonly HashSet<AggregateProjectEntry> _aggregates = new HashSet<AggregateProjectEntry>();

Python/Product/Analysis/Values/BoundBuiltinMethodInfo.cs

@@ -50,7 +50,12 @@ public override string Description {
         }
 
         public override IAnalysisSet Call(Node node, AnalysisUnit unit, IAnalysisSet[] args, NameExpression[] keywordArgNames) {
-            return _method.ReturnTypes.GetInstanceType();
+            // Check if method returns self
+            var returnType = _method.ReturnTypes.GetInstanceType();
+            if (args.Length > 0 && returnType.Split(v => v is BuiltinInstanceInfo biif && biif.PythonType == _method.Function?.DeclaringType, out _, out _)) {
+                return args[0]; //  Return actual self (i.e. derived class)
+            }
+            return returnType;
         }
 
         public override IEnumerable<OverloadResult> Overloads {

Python/Product/VSCode/AnalysisVsc/LanguageServer.cs

@@ -213,13 +213,13 @@ public Task<object> ExecuteCommand(JToken token)
         #region TextDocument
         [JsonRpcMethod("textDocument/didOpen")]
         public Task DidOpenTextDocument(JToken token) {
-            _idleTimeTracker.NotifyUserActivity();
+            _idleTimeTracker?.NotifyUserActivity();
             return _server.DidOpenTextDocument(ToObject<DidOpenTextDocumentParams>(token));
         }
 
         [JsonRpcMethod("textDocument/didChange")]
         public void DidChangeTextDocument(JToken token) {
-            _idleTimeTracker.NotifyUserActivity();
+            _idleTimeTracker?.NotifyUserActivity();
 
             var @params = ToObject<DidChangeTextDocumentParams>(token);
             var version = @params.textDocument.version;
@@ -281,13 +281,13 @@ public Task<TextEdit[]> WillSaveWaitUntilTextDocument(JToken token)
 
         [JsonRpcMethod("textDocument/didSave")]
         public Task DidSaveTextDocument(JToken token) {
-            _idleTimeTracker.NotifyUserActivity();
+            _idleTimeTracker?.NotifyUserActivity();
             return _server.DidSaveTextDocument(ToObject<DidSaveTextDocumentParams>(token));
         }
 
         [JsonRpcMethod("textDocument/didClose")]
         public Task DidCloseTextDocument(JToken token) {
-            _idleTimeTracker.NotifyUserActivity();
+            _idleTimeTracker?.NotifyUserActivity();
             return _server.DidCloseTextDocument(ToObject<DidCloseTextDocumentParams>(token));
         }
         #endregion

Python/Tests/Analysis/LanguageServerTests.cs

@@ -380,6 +380,15 @@ public async Task CompletionInWithStatement() {
             await s.UnloadFileAsync(u);
         }
 
+        [TestMethod, Priority(0)]
+        public async Task CompletionInWithStatementDerivedClass() {
+            using (var s = await CreateServer()) {
+                var u = await AddModule(s, "with open(x) as fs:\n  fs. ");
+                await AssertCompletion(s, u, new[] { "read", "write" }, Enumerable.Empty<string>(), new SourceLocation(2, 6));
+                await s.UnloadFileAsync(u);
+            }
+        }
+
         [TestMethod, Priority(0)]
         public async Task CompletionInImport() {
             var s = await CreateServer();