Merge PR #162

Author: AlexanderSher

src/Analysis/Engine/Test/LineFormatterTests.cs

@@ -208,9 +208,11 @@ public void UnaryOperators(string code, string expected) {
             AssertSingleLineFormat(code, expected);
         }
 
-        [TestMethod, Priority(0)]
-        public void EqualsWithTypeHints() {
-            AssertSingleLineFormat("def foo(x:int=3,x=100.)", "def foo(x: int = 3, x=100.)");
+        [DataRow("def foo(x:int=3,x=100.)", "def foo(x: int = 3, x=100.)")]
+        [DataRow("def foo(x:Union[int,str]=3,x=100.)", "def foo(x: Union[int, str] = 3, x=100.)")]
+        [DataTestMethod, Priority(0)]
+        public void EqualsWithTypeHints(string code, string expected) {
+            AssertSingleLineFormat(code, expected);
         }
 
         [TestMethod, Priority(0)]

src/LanguageServer/Impl/Implementation/LineFormatter.cs

@@ -149,9 +149,35 @@ public TextEdit[] FormatLine(int line) {
                         builder.Append(token);
                         break;
 
-                    case TokenKind.Assign when token.IsInsideFunctionArgs && prev?.PrevNonIgnored?.Kind != TokenKind.Colon:
-                        builder.Append(token);
+                    case TokenKind.Assign when token.IsInsideFunctionArgs:
+                        // Search backwards through the tokens looking for a colon for this argument,
+                        // indicating that there's a type hint and spacing should surround the equals.
+                        for (var p = token.PrevNonIgnored; p != null; p = p.PrevNonIgnored) {
+                            if (p == token.Inside) {
+                                // Hit the surrounding left parenthesis, so stop the search.
+                                builder.Append(token);
+                                break;
+                            }
+
+                            if (p.Inside != token.Inside) {
+                                // Inside another grouping than the =, so skip over it.
+                                continue;
+                            }
+
+                            if (p.Kind == TokenKind.Comma) {
+                                // Found a comma, indicating the end of another argument, so stop.
+                                builder.Append(token);
+                                break;
+                            }
+
+                            if (p.Kind == TokenKind.Colon) {
+                                // Found a colon before hitting another argument or the opening parenthesis, so add spacing.
+                                AppendTokenEnsureWhiteSpacesAround(builder, token);
+                                break;
+                            }
+                        }
                         break;
+
                     case TokenKind.Assign:
                         AppendTokenEnsureWhiteSpacesAround(builder, token);
                         break;