Repository: ilevkivskyi/com2ann

Issue:
When using --wrap-signatures, return annotations with commas also get wrapped
For example:
```python
def func(arg1: SomeVeryLongType1,
         arg2: SomeVeryLongType2) -> Dict[str,
         int]:
    ...
```
This is not a syntax error, but looks ugly.

Discussion:

Patch:
diff --git a/src/com2ann.py b/src/com2ann.py
index b1683e3..166672e 100644
--- a/src/com2ann.py
+++ b/src/com2ann.py
@@ -459,6 +459,7 @@ def wrap_function_header(header: str) -> List[str]:
     parts: List[str] = []
     next_part = ''
     nested = 0
+    complete = False  # Did we split all the arguments inside (...)?
     indent: Optional[int] = None
 
     for i, c in enumerate(header):
@@ -468,7 +469,9 @@ def wrap_function_header(header: str) -> List[str]:
                 indent = i + 1
         if c in ')]}':
             nested -= 1
-        if c == ',' and nested == 1:
+            if not nested:
+                complete = True
+        if c == ',' and nested == 1 and not complete:
             next_part += c
             parts.append(next_part)
             next_part = ''


Test patch:
diff --git a/src/test_com2ann.py b/src/test_com2ann.py
index fafb71c..6cef088 100644
--- a/src/test_com2ann.py
+++ b/src/test_com2ann.py
@@ -411,6 +411,19 @@ class FunctionTestCase(BaseTestCase):
                          *fake_receipts: str) -> None:
                 pass
             """, False, False, 10)
+        self.check(
+            """
+            def embezzle(self, account, funds=MANY, *fake_receipts):
+                # type: (str, int, *str) -> Dict[str, Dict[str, int]]
+                pass
+            """,
+            """
+            def embezzle(self,
+                         account: str,
+                         funds: int = MANY,
+                         *fake_receipts: str) -> Dict[str, Dict[str, int]]:
+                pass
+            """, False, False, 10)
 
     def test_decorator_body(self) -> None:
         self.check(