diff --git a/src/docformatter/format.py b/src/docformatter/format.py index 352aac3..584ff6a 100644 --- a/src/docformatter/format.py +++ b/src/docformatter/format.py @@ -381,16 +381,50 @@ def _get_function_docstring_newlines( # noqa: PLR0911 return 0 -def _get_module_docstring_newlines() -> int: +def _get_module_docstring_newlines( + tokens: list[tokenize.TokenInfo], + index: int, +) -> int: """Return number of newlines after a module docstring. docformatter_8.2: One blank line after a module docstring. + docformatter_8.4: Two blank lines after a module docstring if a top-level + function or class definition follows it (PEP 8: top-level definitions are + surrounded by two blank lines). + + Parameters + ---------- + tokens : list + A list of tokens from the source code. + index : int + The index of the docstring token in the list of tokens. Returns ------- newlines : int The number of newlines to insert after the docstring. """ + j = index + 1 + + while j < len(tokens): + if tokens[j].type in ( + tokenize.NL, + tokenize.NEWLINE, + tokenize.INDENT, + tokenize.DEDENT, + ): + j += 1 + continue + + # A decorated definition is still a definition; skip the decorator line. + if tokens[j].type == tokenize.OP and tokens[j].string == "@": + return 2 + + if _classify.is_definition_line(tokens[j]): + return 2 + + break + return 1 @@ -419,7 +453,7 @@ def _get_newlines_by_type( return 0 elif _classify.is_module_docstring(tokens, index): # print("Module") - return _get_module_docstring_newlines() + return _get_module_docstring_newlines(tokens, index) elif _classify.is_class_docstring(tokens, index): # print("Class") return _get_class_docstring_newlines(tokens, index) diff --git a/tests/_data/string_files/format_functions.toml b/tests/_data/string_files/format_functions.toml index 36e3a9e..a47522e 100644 --- a/tests/_data/string_files/format_functions.toml +++ b/tests/_data/string_files/format_functions.toml @@ -3,20 +3,72 @@ # for creating a TokenInfo() object. [module_docstring_followed_by_string] +source = '''"""Module docstring.""" + +MY_CONSTANT = "a string" +''' expected = 1 [module_docstring_followed_by_code] +source = '''"""Module docstring.""" + +import os +''' expected = 1 [module_docstring_followed_by_comment_then_code] +source = '''"""Module docstring.""" + +# A comment. +import os +''' expected = 1 [module_docstring_followed_by_comment_then_string] +source = '''"""Module docstring.""" + +# A comment. +MY_CONSTANT = "a string" +''' expected = 1 [module_docstring_in_black] +source = '''"""Module docstring.""" + +import os +''' expected = 1 +# PEP 8: top-level function and class definitions are surrounded by two blank +# lines, so a module docstring that is followed by one gets two newlines. +[module_docstring_followed_by_def] +source = '''"""Module docstring.""" + + +def foo(): + """Foo.""" +''' +expected = 2 + +[module_docstring_followed_by_class] +source = '''"""Module docstring.""" + + +class Foo: + """Foo class.""" +''' +expected = 2 + +[module_docstring_followed_by_decorated_def] +source = '''"""Module docstring.""" + + +@decorator +def foo(): + """Foo.""" +''' +expected = 2 + [class_docstring_followed_by_statement] source = ''' class MyClass: diff --git a/tests/formatter/test_format_functions.py b/tests/formatter/test_format_functions.py index 5afd39e..45e789c 100644 --- a/tests/formatter/test_format_functions.py +++ b/tests/formatter/test_format_functions.py @@ -71,12 +71,19 @@ def _get_docstring_token_and_index(tokens): "module_docstring_followed_by_comment_then_code", "module_docstring_followed_by_comment_then_string", "module_docstring_in_black", + "module_docstring_followed_by_def", + "module_docstring_followed_by_class", + "module_docstring_followed_by_decorated_def", ], ) def test_module_docstring_newlines(test_key): + source = TEST_STRINGS[test_key]["source"] expected = TEST_STRINGS[test_key]["expected"] - result = _format._get_module_docstring_newlines() + tokens = _get_tokens(source) + index = _get_docstring_token_and_index(tokens) + + result = _format._get_module_docstring_newlines(tokens, index) assert ( result == expected ), f"\nFailed {test_key}:\nExpected {expected}\nGot {result}"