diff --git a/src/docformatter/format.py b/src/docformatter/format.py index 0f32942..b710023 100644 --- a/src/docformatter/format.py +++ b/src/docformatter/format.py @@ -320,7 +320,9 @@ def _format_code( not in {tokenize.NL, tokenize.NEWLINE} or modified_tokens[-2][1] != ":" and modified_tokens[-2][0] != tokenize.COMMENT - or modified_tokens[-2][4][:3] != "def" + or not modified_tokens[-2][4] + .lstrip() + .startswith(("def", "class")) ): modified_tokens.append( (token_type, token_string, start, end, line) diff --git a/tests/test_format_code.py b/tests/test_format_code.py index 5cd5254..8859012 100644 --- a/tests/test_format_code.py +++ b/tests/test_format_code.py @@ -499,6 +499,80 @@ def foo(): ''' ) + @pytest.mark.unit + @pytest.mark.parametrize("args", [[""]]) + def test_format_code_with_empty_lines_class_docstring( + self, test_args, args + ): + """No blank lines before a class's docstring.""" + uut = Formatter( + test_args, + sys.stderr, + sys.stdin, + sys.stdout, + ) + + assert '''\ +class Foo: + """Hello foo and this is a docstring. + + More stuff. + """ +''' == uut._format_code( + '''\ +class Foo: + """ + Hello + foo and this is a docstring. + + More stuff. + """ +''' + ) + assert '''\ +def foo(): + class Foo: + """Summary.""" + pass +''' == uut._format_code( + '''\ +def foo(): + class Foo: + + """Summary.""" + pass +''' + ) + + @pytest.mark.unit + @pytest.mark.parametrize("args", [[""]]) + def test_format_code_with_empty_lines_method_docstring( + self, test_args, args + ): + """No blank lines before a method's docstring.""" + uut = Formatter( + test_args, + sys.stderr, + sys.stdin, + sys.stdout, + ) + + assert '''\ +class Foo: + def foo(self): + """Summary.""" + pass +''' == uut._format_code( + '''\ +class Foo: + def foo(self): + + + """Summary.""" + pass +''' + ) + @pytest.mark.unit @pytest.mark.parametrize("args", [[""]]) def test_format_code_with_trailing_whitespace(self, test_args, args): @@ -860,13 +934,15 @@ class TestClass: ..py.method: big_method() """ ''' - assert docstring == uut._do_format_code('''\ + assert docstring == uut._do_format_code( + '''\ class TestClass: """This is a class docstring. :cvar test_int: a class attribute. ..py.method: big_method() """ -''') +''' + ) class TestFormatCodeRanges: