Merge pull request #50 from BrianPugh/fix-multiline-inspect

Fix multiline string inspect

.flake8

@@ -21,6 +21,7 @@ exclude =
     .venv
     examples/
     belay/snippets/
+    tests/test_inspect.py
 
 extend-immutable-calls =
     Argument

belay/inspect.py

@@ -1,16 +1,63 @@
 import inspect
 import re
+from io import StringIO
+from tokenize import (
+    COMMENT,
+    DEDENT,
+    INDENT,
+    NEWLINE,
+    OP,
+    STRING,
+    generate_tokens,
+    untokenize,
+)
 from typing import Tuple
 
 _pat_no_decorators = re.compile(
     r"^(\s*def\s)|(\s*async\s+def\s)|(.*(?<!\w)lambda(:|\s))"
 )
 
 
+class _NoAction(Exception):
+    pass
+
+
+def _dedent_tokenizer(code):
+    indent_to_remove = ""
+    for (
+        token_type,
+        string,
+        (start_line, start_col),
+        (end_line, end_col),  # noqa: B007
+        _,
+    ) in generate_tokens(StringIO(code).readline):
+        print(f"{token_type=} {string=} {start_line=} {end_line=}")
+        if start_line == 1 and start_col == 0:
+            # First Token
+            if token_type != INDENT:
+                # No action to perform
+                raise _NoAction
+            indent_to_remove = string
+        if start_col == 0:
+            if token_type == INDENT:
+                if not string.startswith(indent_to_remove):
+                    raise IndentationError
+                string = string[len(indent_to_remove) :]
+        yield token_type, string
+
+
+def _dedent(code):
+    try:
+        return untokenize(_dedent_tokenizer(code))
+    except _NoAction:
+        return code
+
+
 def getsource(f) -> Tuple[str, int, str]:
-    """Get source code data without decorators.
+    """Get source code with mild post processing.
 
-    Trims leading whitespace and removes decorators.
+    * strips leading decorators.
+    * Trims leading whitespace indent.
 
     Parameters
     ----------
@@ -20,7 +67,7 @@ def getsource(f) -> Tuple[str, int, str]:
     Returns
     -------
     src_code: str
-        Source code.
+        Source code. Always ends in a newline.
     src_lineno: int
         Line number of code begin.
     src_file: str
@@ -39,11 +86,9 @@ def getsource(f) -> Tuple[str, int, str]:
 
     lines = lines[offset:]
 
-    # Trim leading whitespace
-    n_leading_whitespace = len(lines[0]) - len(lines[0].lstrip())
-    lines = [line[n_leading_whitespace:] for line in lines]
-
     src_code = "".join(lines)
     src_lineno += offset
 
+    src_code = _dedent(src_code)
+
     return src_code, src_lineno, src_file

tests/test_inspect.py

@@ -1,3 +1,10 @@
+"""Tests extensions to the builtin inspect module.
+
+Note: ``untokenize`` doesn't properly preserve inter-token spacing, so this test vector
+may need to change while still remaining valid.
+"""
+
+
 from importlib.machinery import SourceFileLoader
 
 import pytest
@@ -73,7 +80,7 @@ def test_getsource_decorated_4(foo):
 def test_getsource_decorated_5(foo):
     """Removes leading indent."""
     code, lineno, file = belay.inspect.getsource(foo.foo_decorated_5)
-    assert code == "def foo_decorated_5(arg1, arg2):\n    return arg1 + arg2\n"
+    assert code == "def foo_decorated_5 (arg1 ,arg2 ):\n    return arg1 +arg2 \n"
     assert lineno == 45
     assert file == foo.__file__
 
@@ -84,3 +91,65 @@ def test_getsource_decorated_6(foo):
     assert code == "def foo_decorated_6(arg1, arg2):\n    return arg1 + arg2\n"
     assert lineno == 51
     assert file == foo.__file__
+
+
+def test_getsource_decorated_7(foo):
+    """Double decorated."""
+    code, lineno, file = belay.inspect.getsource(foo.foo_decorated_7)
+    assert (
+        code
+        == 'def foo_decorated_7(arg1, arg2):\n    return """This\n    is\na\n  multiline\n             string.\n"""\n'
+    )
+    assert lineno == 56
+    assert file == foo.__file__
+
+
+def test_getsource_nested():
+    def foo():
+        bar = 5
+        return 7
+
+    code, lineno, file = belay.inspect.getsource(foo)
+    assert code == "def foo ():\n    bar =5 \n    return 7 \n"
+    assert file == __file__
+
+
+def test_getsource_nested_multiline_string():
+    for _ in range(1):
+
+        def foo(arg1, arg2):
+            return """This
+    is
+a
+  multiline
+             string.
+"""
+
+    code, lineno, file = belay.inspect.getsource(foo)
+    assert (
+        code
+        == 'def foo (arg1 ,arg2 ):\n    return """This\n    is\na\n  multiline\n             string.\n"""\n'
+    )
+    assert file == __file__
+
+
+def test_getsource_nested_multiline_string():
+    # fmt: off
+    def bar(a, b):
+        return a * b
+
+    for _ in range(1):
+
+        def foo(arg1, arg2):
+            return bar(
+arg1,
+    arg2
+)
+
+    # fmt: on
+    code, lineno, file = belay.inspect.getsource(foo)
+    assert (
+        code
+        == "def foo (arg1 ,arg2 ):\n    return bar (\n    arg1 ,\n    arg2 \n    )\n"
+    )
+    assert file == __file__

tests/test_inspect/foo.py

@@ -50,3 +50,13 @@ def foo_decorated_5(arg1, arg2):
 @decorator
 def foo_decorated_6(arg1, arg2):
     return arg1 + arg2
+
+
+@decorator
+def foo_decorated_7(arg1, arg2):
+    return """This
+    is
+a
+  multiline
+             string.
+"""