Repository: heavenshell/py-doq

Issue:
Multi-line return annotations
Hi,

Thanks for the super useful tool! Just wanted to report one minor/non-urgent issue that I've been running into recently --

Specifically, running `doq` on:
```python
from typing import Tuple


def x(a, b, c) -> Tuple[
    int,
    int,
    int,
]:
    pass
```

produces:
```python
from typing import Tuple


def x(a, b, c) -> Tuple[
    """x.

    :param a:
    :param b:
    :param c:
    """
    int,
    int,
    int,
]:
    pass
```

Seems like two issues:
- `detect_insert_point()` breaks in outputter.py
- `parse_return_type()` breaks in parser.py; note that there are no return annotations above

For the latter: I think `parse_return_type()` could be replaced by checking the parso tree, for example in `parse_defs()`:
```python
        if d.children[3].value == "->":
            # Get return-type annotation
            return_type = d.children[4].get_code().strip() 
        else:
            # No annotation
            return_type = None
```
(child indices are from [here](https://github.com/davidhalter/parso/blob/034a9e89443261e7f434fcb1fbb807b14991b766/parso/python/tree.py#L536))

Some thought may also need to be put into how indents, etc are handled?

@heavenshell if you don't have time to fix, I might spontaneously submit a PR sometime next month. :slightly_smiling_face: 

Comments:

Patch:
diff --git a/doq/outputter.py b/doq/outputter.py
index 4099b06..8fe6de7 100644
--- a/doq/outputter.py
+++ b/doq/outputter.py
@@ -14,13 +14,20 @@ class StringOutptter(BaseOutputter):
             return start
         else:
             for i, line in enumerate(lines[start:end]):
-                # Found end of signature without type
                 if line.endswith('):'):
+                    # Found end of signature without type
                     return start + i + 1
                 elif re.search(r'\):', line):
                     return start + i + 1
-                # Found end of signature with type
+                elif re.search(r'\]:', line):
+                    # Found end of signature type
+                    #   def foo(a, b) -> Tuple[
+                    #       int,
+                    #   ]:
+                    #       pass
+                    return start + i + 1
                 elif re.search(r'->(.*):', line):
+                    # Found end of signature with type
                     return start + i + 1
 
         return start
diff --git a/doq/parser.py b/doq/parser.py
index 34805a7..70755bc 100644
--- a/doq/parser.py
+++ b/doq/parser.py
@@ -88,14 +88,9 @@ def parse_defs(module, omissions=None, ignore_exception=False, ignore_yield=Fals
 
             params.append(arguments)
 
-        # parso does not have return type. So parse from signature.
-        next_node = d.get_suite().get_first_leaf().get_next_sibling()
-        stmt_start_lineno = next_node.start_pos[0] if next_node else 2
-        return_type = parse_return_type(
-            code=code,
-            start_lineno=start_lineno,
-            end_lineno=stmt_start_lineno - 1,
-        )
+        return_type = None
+        if d.children[3].value == '->':
+            return_type = d.children[4].get_code().strip()
 
         yields = []
         if ignore_yield is False:


Test patch:
diff --git a/tests/test_outputter.py b/tests/test_outputter.py
index 6229985..5368858 100644
--- a/tests/test_outputter.py
+++ b/tests/test_outputter.py
@@ -107,6 +107,40 @@ class StringOutptterTestCase(TestCase):
         ])
         self.assertEqual(expected, output)
 
+    def test_multi_return_type(self):
+        lines = [
+            'def foo(arg1) -> List[',
+            '    int,',
+            '    int,',
+            ']:',
+            '    pass',
+        ]
+        docstrings = [{
+            'docstring': '"""foo.\n\n:param arg1:\n:rtype List[\n    int,\n    int,\n]:\n"""',
+            'start_lineno': 1,
+            'start_col': 0,
+            'end_lineno': 7,
+            'end_col': 0,
+            'is_doc_exists': False,
+        }]
+        output = StringOutptter().format(lines=lines, docstrings=docstrings, indent=4)
+        expected = '\n'.join([
+            'def foo(arg1) -> List[',
+            '    int,',
+            '    int,',
+            ']:',
+            '    """foo.',
+            '',
+            '    :param arg1:',
+            '    :rtype List[',
+            '        int,',
+            '        int,',
+            '    ]:',
+            '    """',
+            '    pass',
+        ])
+        self.assertEqual(expected, output)
+
 
 class JSONOutptterTestCase(TestCase):
     def test_same_lines(self):
diff --git a/tests/test_parser.py b/tests/test_parser.py
index 404112a..261b8e4 100644
--- a/tests/test_parser.py
+++ b/tests/test_parser.py
@@ -484,6 +484,37 @@ class ParseTestCase(TestCase):
             actual,
         )
 
+    def test_with_return_multi_line_type(self):
+        line = '\n'.join([
+            'def foo(arg1) -> Tuple[',
+            '    int,',
+            '    int,',
+            ']:',
+            '    pass',
+        ])
+        actual = parse(line)[0]
+        self.assertDictEqual(
+            {
+                'name': 'foo',
+                'params': [
+                    {
+                        'argument': 'arg1',
+                        'annotation': None,
+                        'default': None,
+                    },
+                ],
+                'return_type': 'Tuple[\n    int,\n    int,\n]',
+                'start_lineno': 1,
+                'start_col': 0,
+                'end_lineno': 5,
+                'end_col': 8,
+                'is_doc_exists': False,
+                'exceptions': [],
+                'yields': [],
+            },
+            actual,
+        )
+
     def test_with_defs(self):
         line = '\n'.join([
             'def bar(arg1) -> List[str]:',