Allow standalone comments to close code blocks

Fixes #16
Fixes #32

diff --git a/README.md b/README.md
index 9685172d7bf29fba584701b1a0bae670def97278..b380ee7b90c43f6d8ff017fccb2247be6b0db489 100644 (file)
--- a/README.md
+++ b/README.md
@@ -518,6 +518,8 @@ More details can be found in [CONTRIBUTING](CONTRIBUTING.md).
 * generalized star expression handling, including double stars; this
   fixes multiplication making expressions "unsafe" for trailing commas (#132)
 
+* fixed comment indentation when a standalone comment closes a block (#16, #32)
+
 * fixed `--diff` not showing entire path (#130)
 
 * fixed parsing of complex expressions after star and double stars in

diff --git a/black.py b/black.py
index da645a1fad535382255cbd0801acd89991d4c601..15a754742cc2212656d324fab4422ddad6982289 100644 (file)
--- a/black.py
+++ b/black.py
@@ -1158,7 +1158,16 @@ class LineGenerator(Visitor[Line]):
 
     def visit_DEDENT(self, node: Node) -> Iterator[Line]:
         """Decrease indentation level, maybe yield a line."""
-        # DEDENT has no value. Additionally, in blib2to3 it never holds comments.
+        # The current line might still wait for trailing comments.  At DEDENT time
+        # there won't be any (they would be prefixes on the preceding NEWLINE).
+        # Emit the line then.
+        yield from self.line()
+
+        # While DEDENT has no value, its prefix may contain standalone comments
+        # that belong to the current indentation level.  Get 'em.
+        yield from self.visit_default(node)
+
+        # Finally, emit the dedent.
         yield from self.line(-1)
 
     def visit_stmt(

diff --git a/blib2to3/pgen2/driver.py b/blib2to3/pgen2/driver.py
index 5cdd2e59e6156cecd8ffa3cebc822b4be20ef7ae..2265c26443d2f45280112653f5f25eda59360b07 100644 (file)
--- a/blib2to3/pgen2/driver.py
+++ b/blib2to3/pgen2/driver.py
@@ -43,6 +43,7 @@ class Driver(object):
         p.setup()
         lineno = 1
         column = 0
+        indent_columns = []
         type = value = start = end = line_text = None
         prefix = ""
         for quintuple in tokens:
@@ -72,12 +73,16 @@ class Driver(object):
             if type in {token.INDENT, token.DEDENT}:
                 _prefix = prefix
                 prefix = ""
+            if type == token.DEDENT:
+                _indent_col = indent_columns.pop()
+                prefix, _prefix = self._partially_consume_prefix(_prefix, _indent_col)
             if p.addtoken(type, value, (prefix, start)):
                 if debug:
                     self.logger.debug("Stop.")
                 break
             prefix = ""
             if type == token.INDENT:
+                indent_columns.append(len(value))
                 if _prefix.startswith(value):
                     # Don't double-indent.  Since we're delaying the prefix that
                     # would normally belong to INDENT, we need to put the value
@@ -114,6 +119,35 @@ class Driver(object):
         tokens = tokenize.generate_tokens(io.StringIO(text).readline)
         return self.parse_tokens(tokens, debug)
 
+    def _partially_consume_prefix(self, prefix, column):
+        lines = []
+        current_line = ""
+        current_column = 0
+        wait_for_nl = False
+        for char in prefix:
+            current_line += char
+            if wait_for_nl:
+                if char == '\n':
+                    if current_line.strip() and current_column < column:
+                        res = ''.join(lines)
+                        return res, prefix[len(res):]
+
+                    lines.append(current_line)
+                    current_line = ""
+                    current_column = 0
+                    wait_for_nl = False
+            elif char == ' ':
+                current_column += 1
+            elif char == '\t':
+                current_column += 4
+            elif char == '\n':
+                # enexpected empty line
+                current_column = 0
+            else:
+                # indent is finished
+                wait_for_nl = True
+        return ''.join(lines), current_line
+
 
 def _generate_pickle_name(gt):
     head, tail = os.path.splitext(gt)

diff --git a/tests/comments2.py b/tests/comments2.py
index 848ddb1ae8541a9dcf25b429e6481fe5da204cf4..73fff32a3b933d6629aee64650099177c35062ac 100644 (file)
--- a/tests/comments2.py
+++ b/tests/comments2.py
@@ -158,7 +158,7 @@ else:
     # for compiler in compilers.values():
     # add_compiler(compiler)
     add_compiler(compilers[(7.0, 32)])
-# add_compiler(compilers[(7.1, 64)])
+    # add_compiler(compilers[(7.1, 64)])
 
 # Comment before function.
 
@@ -238,8 +238,8 @@ short
         if False:
             continue
 
-    # and round and round we go
-    # and round and round we go
+            # and round and round we go
+        # and round and round we go
 
     # let's return
     return Node(

diff --git a/tests/comments5.py b/tests/comments5.py
new file mode 100644 (file)
index 0000000..703922d
--- /dev/null
+++ b/tests/comments5.py
@@ -0,0 +1,31 @@
+while True:
+    if something.changed:
+        do.stuff()  # trailing comment
+        # Comment belongs to the `if` block.
+    # This one belongs to the `while` block.
+
+    # Should this one, too?  I guess so.
+
+# This one is properly standalone now.
+
+for i in range(100):
+    # first we do this
+    if i % 33 == 0:
+        break
+
+    # then we do this
+    print(i)
+    # and finally we loop around
+
+with open(some_temp_file) as f:
+    data = f.read()
+
+try:
+    with open(some_other_file) as w:
+        w.write(data)
+
+except OSError:
+    print("problems")
+
+if __name__ == "__main__":
+    main()

diff --git a/tests/test_black.py b/tests/test_black.py
index 6f0ffa364ac749e770e67549ee6e4420bfc8e739..dd3beed99abe0e5b6d7b3f37188fa41420f38493 100644 (file)
--- a/tests/test_black.py
+++ b/tests/test_black.py
@@ -262,6 +262,14 @@ class BlackTestCase(unittest.TestCase):
         black.assert_equivalent(source, actual)
         black.assert_stable(source, actual, line_length=ll)
 
+    @patch("black.dump_to_file", dump_to_stderr)
+    def test_comments5(self) -> None:
+        source, expected = read_data("comments5")
+        actual = fs(source)
+        self.assertFormatEqual(expected, actual)
+        black.assert_equivalent(source, actual)
+        black.assert_stable(source, actual, line_length=ll)
+
     @patch("black.dump_to_file", dump_to_stderr)
     def test_cantfit(self) -> None:
         source, expected = read_data("cantfit")