Comment out With blocks with TODO instead of removing them.

rubberduck-vba · Oct 28, 2018 · d122208 · d122208
1 parent ca8beff
commit d122208
Show file tree

Hide file tree

Showing 2 changed files with 108 additions and 1 deletion.
diff --git a/Rubberduck.CodeAnalysis/QuickFixes/RemoveUnassignedVariableUsageQuickFix.cs b/Rubberduck.CodeAnalysis/QuickFixes/RemoveUnassignedVariableUsageQuickFix.cs
@@ -1,3 +1,7 @@
+using System;
+using System.Collections.Generic;
+using System.Linq;
+using System.ServiceModel;
 using Antlr4.Runtime;
 using Rubberduck.Inspections.Abstract;
 using Rubberduck.Inspections.Concrete;
@@ -23,8 +27,26 @@ public override void Fix(IInspectionResult result)
         {
             var rewriter = _state.GetRewriter(result.QualifiedSelection.QualifiedName);
 
+            if (result.Context.Parent.Parent is VBAParser.WithStmtContext withContext)
+            {
+                var lines = withContext.GetText().Replace("\r", string.Empty).Split('\n');
+                // Assume that the End With is at the appropriate indentation level for the block. Note that this could
+                // over-indent or under-indent some lines if statement separators are being used, but meh. 
+                var padding = new string(' ', lines.Last().IndexOf(Tokens.End, StringComparison.Ordinal));
+
+                var replacement = new List<string>
+                {
+                    $"{Tokens.CommentMarker}TODO - {result.Description}",
+                    $"{Tokens.CommentMarker}{padding}{lines.First()}"
+                };
+                replacement.AddRange(lines.Skip(1)
+                    .Select(line => Tokens.CommentMarker + line));
+
+                rewriter.Replace(withContext, string.Join(Environment.NewLine, replacement));
+                return;
+            }
             var assignmentContext = result.Context.GetAncestor<VBAParser.LetStmtContext>() ??
-                                                  (ParserRuleContext)result.Context.GetAncestor<VBAParser.CallStmtContext>();
+                                    (ParserRuleContext)result.Context.GetAncestor<VBAParser.CallStmtContext>();
 
             rewriter.Remove(assignmentContext);
         }

diff --git a/RubberduckTests/QuickFixes/RemoveUnassignedVariableUsageQuickFixTests.cs b/RubberduckTests/QuickFixes/RemoveUnassignedVariableUsageQuickFixTests.cs
@@ -39,5 +39,90 @@ public void UnassignedVariableUsage_QuickFixWorks()
                 Assert.AreEqual(expectedCode, state.GetRewriter(component).GetText());
             }
         }
+
+        // See https://github.com/rubberduck-vba/Rubberduck/issues/3636
+        [Test]
+        [Category("QuickFixes")]
+        public void UnassignedVariableUsage_QuickFixWorksWithBlock()
+        {
+            const string inputCode =
+                @"Sub test()
+    Dim wb As Workbook
+    With wb
+        Debug.Print .Name
+        Debug.Print .Name
+        Debug.Print .Name
+    End With
+End Sub";
+
+            const string expectedCode =
+                @"Sub test()
+    Dim wb As Workbook
+    'TODO - {0}
+'    With wb
+'        Debug.Print .Name
+'        Debug.Print .Name
+'        Debug.Print .Name
+'    End With
+End Sub";
+
+            var vbe = MockVbeBuilder.BuildFromSingleStandardModule(inputCode, out var component);
+            using (var state = MockParser.CreateAndParse(vbe.Object))
+            {
+                var inspection = new UnassignedVariableUsageInspection(state);
+                var inspectionResult = inspection.GetInspectionResults(CancellationToken.None).First();
+                var expected = string.Format(expectedCode, inspectionResult.Description);
+
+                new RemoveUnassignedVariableUsageQuickFix(state).Fix(inspectionResult);
+                var actual = state.GetRewriter(component).GetText();
+                Assert.AreEqual(expected, actual);
+            }
+        }
+
+        [Test]
+        [Category("QuickFixes")]
+        public void UnassignedVariableUsage_QuickFixWorksNestedWithBlock()
+        {
+            const string inputCode =
+                @"Sub test()
+    Dim wb As Workbook
+    Dim ws As Worksheet
+    With wb
+        Debug.Print .Name
+        With ws
+            Debug.Print .Name
+            Debug.Print .Name
+            Debug.Print .Name
+        End With
+    End With
+End Sub";
+
+            const string expectedCode =
+                @"Sub test()
+    Dim wb As Workbook
+    Dim ws As Worksheet
+    With wb
+        Debug.Print .Name
+        'TODO - {0}
+'        With ws
+'            Debug.Print .Name
+'            Debug.Print .Name
+'            Debug.Print .Name
+'        End With
+    End With
+End Sub";
+
+            var vbe = MockVbeBuilder.BuildFromSingleStandardModule(inputCode, out var component);
+            using (var state = MockParser.CreateAndParse(vbe.Object))
+            {
+                var inspection = new UnassignedVariableUsageInspection(state);
+                var inspectionResult = inspection.GetInspectionResults(CancellationToken.None).First();
+                var expected = string.Format(expectedCode, inspectionResult.Description);
+
+                new RemoveUnassignedVariableUsageQuickFix(state).Fix(inspectionResult);
+                var actual = state.GetRewriter(component).GetText();
+                Assert.AreEqual(expected, actual);
+            }
+        }
     }
 }