Make a minor adjustment to the whitespace stripper. (#31)

Author: lukesneeringer

packages/gapic-generator/gapic/generator/formatter.py

@@ -32,10 +32,10 @@ def fix_whitespace(code: str) -> str:
     code = re.sub(r'[ ]+\n', '\n', code)
 
     # Ensure at most two blank lines before top level definitions.
-    code = re.sub(r'\s+\n\n\n(class|def|@)', r'\n\n\n\1', code)
+    code = re.sub(r'\s+\n\s*\n\s*\n(class|def|@|#)', r'\n\n\n\1', code)
 
     # Ensure at most one line before nested definitions.
-    code = re.sub(r'\s+\n\n(    )+(class|def|@)', r'\n\n\1\2', code)
+    code = re.sub(r'\s+\n\s*\n(    )+(class|def|@|#)', r'\n\n\1\2', code)
 
     # All files shall end in one and exactly one line break.
     return f'{code.rstrip()}\n'

packages/gapic-generator/tests/unit/generator/test_formatter.py

@@ -88,5 +88,44 @@ def too_far_down(self):
     """)
 
 
+def test_fix_whitespace_intermediate_whitespace():
+    assert formatter.fix_whitespace(textwrap.dedent("""\
+    class JustAClass:
+        def foo(self):
+            pass
+        \
+
+
+        @property
+        def too_far_down(self):
+            return 42
+    """)) == textwrap.dedent("""\
+    class JustAClass:
+        def foo(self):
+            pass
+
+        @property
+        def too_far_down(self):
+            return 42
+    """)
+
+
+def test_fix_whitespace_comment():
+    assert formatter.fix_whitespace(textwrap.dedent("""\
+    def do_something():
+        do_first_thing()
+
+
+        # Something something something.
+        do_second_thing()
+    """)) == textwrap.dedent("""\
+    def do_something():
+        do_first_thing()
+
+        # Something something something.
+        do_second_thing()
+    """)
+
+
 def test_file_newline_ending():
     assert formatter.fix_whitespace('') == '\n'