Correctly handle multi-line imports.

Author: domdfcoding

.pre-commit-config.yaml

@@ -28,6 +28,12 @@ repos:
       args: [--allow-git]
     - id: check-docstring-first
 
+ - repo: https://github.com/domdfcoding/flake8-dunder-all
+   rev: v0.0.3
+   hooks:
+    - id: ensure-dunder-all
+      files: ^flake8_dunder_all/.*\.py$
+
  - repo: https://github.com/pre-commit/pygrep-hooks
    rev: v1.5.1
    hooks:

doc-source/requirements.txt

@@ -1,3 +1,4 @@
+autodocsumm>=0.2.0
 default_values>=0.0.7
 domdf_sphinx_theme>=0.0.11
 extras_require
@@ -11,4 +12,3 @@ sphinxcontrib-httpdomain>=1.7.0
 sphinxemoji>=0.1.6
 toctree_plus>=0.0.2
 git+git://github.com/domdfcoding/sphinx-autodoc-typehints.git@typevar-as-pydata
-git+git://github.com/Chilipp/autodocsumm@0a1b6515ba83deb70eb2c356acf956f448536c90

flake8_dunder_all/__init__.py

@@ -31,15 +31,16 @@
 
 # stdlib
 import ast
-from typing import Any, Generator, List, Tuple, Type
+import sys
+from typing import Any, Generator, List, Tuple, Type, Union
 
 # 3rd party
 from domdf_python_tools.paths import PathPlus
 from domdf_python_tools.terminal_colours import Fore
 from domdf_python_tools.utils import stderr_writer
 
 # this package
-from flake8_dunder_all.utils import get_docstring_lineno
+from flake8_dunder_all.utils import get_docstring_lineno, mark_text_ranges
 
 __author__: str = "Dominic Davis-Foster"
 __copyright__: str = "2020 Dominic Davis-Foster"
@@ -55,16 +56,21 @@
 class Visitor(ast.NodeVisitor):
 	"""
 	AST :class:`~ast.NodeVisitor` to check a module has defined ``__all__``, and add one if it not.
+
+	:param use_endlineno: Flag to indicate whether the end_lineno functionality is available.
+		This functionality is available on Python 3.8 and above, or when the tree has been passed through
+		:func:`flake8_dunder_all.utils.mark_text_ranges``.
 	"""
 
-	def __init__(self) -> None:
-		self.found_all = False
+	found_all: bool  #: Flag to indicate a ``__all__`` variable has been found in the AST.
+	last_import: int  #: The lineno of the last top-level import
+	members: List[str]  #: List of functions and classed defined in the AST
 
-		# List of functions and classed defined in this module
+	def __init__(self, use_endlineno: bool = False) -> None:
+		self.found_all = False
 		self.members: List[str] = []
-
-		# Lineno of last top-level import
 		self.last_import = 0
+		self.use_endlineno = use_endlineno
 
 	def visit_Name(self, node: ast.Name):
 		"""
@@ -78,6 +84,16 @@ def visit_Name(self, node: ast.Name):
 		else:
 			self.generic_visit(node)
 
+	def handle_def(self, node: Union[ast.FunctionDef, ast.AsyncFunctionDef, ast.ClassDef]):
+		"""
+		Handles ``def foo(): ...``, ``async def foo(): ...`` and ``class Foo: ...``.
+
+		:param node: The node being visited.
+		"""
+
+		if not node.name.startswith("_"):
+			self.members.append(node.name)
+
 	def visit_FunctionDef(self, node: ast.FunctionDef):
 		"""
 		Visit ``def foo(): ...``.
@@ -86,8 +102,7 @@ def visit_FunctionDef(self, node: ast.FunctionDef):
 		"""
 
 		# Don't generic visit
-		if not node.name.startswith("_"):
-			self.members.append(node.name)
+		self.handle_def(node)
 
 	def visit_AsyncFunctionDef(self, node: ast.AsyncFunctionDef):
 		"""
@@ -97,8 +112,7 @@ def visit_AsyncFunctionDef(self, node: ast.AsyncFunctionDef):
 		"""
 
 		# Don't generic visit
-		if not node.name.startswith("_"):
-			self.members.append(node.name)
+		self.handle_def(node)
 
 	def visit_ClassDef(self, node: ast.ClassDef):
 		"""
@@ -108,8 +122,21 @@ def visit_ClassDef(self, node: ast.ClassDef):
 		"""
 
 		# Don't generic visit
-		if not node.name.startswith("_"):
-			self.members.append(node.name)
+		self.handle_def(node)
+
+	def handle_import(self, node: Union[ast.Import, ast.ImportFrom]):
+		"""
+		Handles ``import foo`` and ``from foo import bar``.
+
+		:param node: The node being visited
+		"""
+
+		if self.use_endlineno:
+			if not node.col_offset and node.end_lineno > self.last_import:  # type: ignore
+				self.last_import = node.end_lineno  # type: ignore
+		else:
+			if not node.col_offset and node.lineno > self.last_import:
+				self.last_import = node.lineno
 
 	def visit_Import(self, node: ast.Import):
 		"""
@@ -119,8 +146,7 @@ def visit_Import(self, node: ast.Import):
 		"""
 
 		# Don't generic visit
-		if not node.col_offset and node.lineno > self.last_import:
-			self.last_import = node.lineno
+		self.handle_import(node)
 
 	def visit_ImportFrom(self, node: ast.ImportFrom):
 		"""
@@ -130,8 +156,7 @@ def visit_ImportFrom(self, node: ast.ImportFrom):
 		"""
 
 		# Don't generic visit
-		if not node.col_offset and node.lineno > self.last_import:
-			self.last_import = node.lineno
+		self.handle_import(node)
 
 
 class Plugin:
@@ -188,12 +213,15 @@ def check_and_add_all(filename: PathPlus, quote_type: str = '"') -> int:
 	filename = PathPlus(filename)
 
 	try:
-		tree = ast.parse(filename.read_text())
+		source = filename.read_text()
+		tree = ast.parse(source)
+		if sys.version_info < (3, 8):  # pragma: no cover (<py38)
+			mark_text_ranges(tree, source)
 	except SyntaxError:
 		stderr_writer(Fore.RED(f"'{filename}' does not appear to be a valid Python source file."))
 		return 4
 
-	visitor = Visitor()
+	visitor = Visitor(use_endlineno=True)
 	visitor.visit(tree)
 
 	if visitor.found_all:

flake8_dunder_all/__main__.py

@@ -51,7 +51,8 @@ def main(argv: Optional[Sequence[str]] = None) -> int:
 
 	"""
 	parser = argparse.ArgumentParser(
-			description=tidy_docstring(main.__doc__), formatter_class=argparse.RawTextHelpFormatter
+			description=tidy_docstring(main.__doc__),
+			formatter_class=argparse.RawTextHelpFormatter,
 			)
 	parser.add_argument('filenames', type=str, nargs='*', help="The filename(s) to lint.", metavar="FILENAME")
 	parser.add_argument(

flake8_dunder_all/utils.py

@@ -32,14 +32,21 @@
 #  Copyright © 1995-2000 Corporation for National Research Initiatives. All rights reserved.
 #  Copyright © 1991-1995 Stichting Mathematisch Centrum. All rights reserved.
 #
+#  mark_text_ranges from Thonny
+#  https://github.com/thonny/thonny/blob/master/thonny/ast_utils.py
+#  Copyright (c) 2020 Aivar Annamaa
+#  MIT Licensed
 
 # stdlib
 import ast
 import re
-from textwrap import dedent, indent
+from textwrap import dedent
 from typing import Optional, Union
 
-__all__ = ["get_docstring_lineno", "tidy_docstring"]
+# 3rd party
+from asttokens.asttokens import ASTTokens
+
+__all__ = ["get_docstring_lineno", "tidy_docstring", "mark_text_ranges"]
 
 
 def get_docstring_lineno(node: Union[ast.FunctionDef, ast.ClassDef, ast.Module]) -> Optional[int]:
@@ -76,3 +83,24 @@ def tidy_docstring(docstring: Optional[str]) -> str:
 	docstring = re.sub("``([^`]*)``", r"'\1'", docstring)
 
 	return f"\n{docstring}"
+
+
+def mark_text_ranges(node: ast.AST, source: str):
+	"""
+	Node is an AST, source is corresponding source as string.
+	Function adds recursively attributes end_lineno and end_col_offset to each node
+	which has attributes lineno and col_offset.
+
+	:param node:
+	:param source: The corresponding source code for the node.
+	"""
+
+	ASTTokens(source, tree=node)
+
+	for child in ast.walk(node):
+		if hasattr(child, "last_token"):
+			child.end_lineno, child.end_col_offset = child.last_token.end  # type: ignore
+
+			if hasattr(child, "lineno"):
+				# Fixes problems with some nodes like binop
+				child.lineno, child.col_offset = child.first_token.start  # type: ignore

requirements.txt

@@ -1,2 +1,3 @@
+asttokens>=1.1
 domdf_python_tools>=0.4.10
 flake8>=3.7

tests/common.py

@@ -78,6 +78,25 @@ async def a_function(): ...
 
 
 
+async def a_function(): ...
+'''
+
+testing_source_j = '''
+from tests.common import (
+		mangled_source,
+		results,
+		testing_source_a,
+		testing_source_b,
+		testing_source_c,
+		testing_source_d,
+		testing_source_e,
+		testing_source_f,
+		testing_source_g,
+		testing_source_h,
+		testing_source_i
+		)
+
+
 async def a_function(): ...
 '''
 

tests/test_flake8_dunder_all.py

@@ -1,6 +1,7 @@
 # stdlib
 import ast
 import re
+import sys
 from typing import List
 
 # 3rd party
@@ -9,7 +10,7 @@
 from domdf_python_tools.terminal_colours import Fore
 
 # this package
-from flake8_dunder_all import Visitor, check_and_add_all
+from flake8_dunder_all import Visitor, check_and_add_all, mark_text_ranges
 from tests.common import (
 		mangled_source,
 		results,
@@ -21,7 +22,8 @@
 		testing_source_f,
 		testing_source_g,
 		testing_source_h,
-		testing_source_i
+		testing_source_i,
+		testing_source_j
 		)
 
 
@@ -44,6 +46,7 @@
 						),
 				pytest.param(testing_source_h, set(), id="from import"),
 				pytest.param(testing_source_i, {'0:0: DALL000 Module lacks __all__.'}, id="lots of lines"),
+				pytest.param(testing_source_j, {'0:0: DALL000 Module lacks __all__.'}, id="multiline import"),
 				]
 		)
 def test_plugin(source, expects):
@@ -73,6 +76,7 @@ def test_plugin(source, expects):
 				pytest.param(testing_source_g, ["a_function"], False, 3, id="async function no __all__"),
 				pytest.param(testing_source_h, [], False, 1, id="from import"),
 				pytest.param(testing_source_i, ["a_function"], False, 3, id="lots of lines"),
+				pytest.param(testing_source_j, ["a_function"], False, 2, id="multiline import"),
 				]
 		)
 def test_visitor(source, members, found_all, last_import):
@@ -84,6 +88,43 @@ def test_visitor(source, members, found_all, last_import):
 	assert visitor.last_import is last_import
 
 
+@pytest.mark.parametrize(
+		"source, members, found_all, last_import",
+		[
+				pytest.param('import foo', [], False, 1, id="just an import"),
+				pytest.param('"""a docstring"""', [], False, 0, id="just a docstring"),
+				pytest.param(testing_source_a, [], False, 3, id="import and docstring"),
+				pytest.param(testing_source_b, ['a_function'], False, 3, id="function no __all__"),
+				pytest.param(testing_source_c, ['Foo'], False, 3, id="class no __all__"),
+				pytest.param(
+						testing_source_d, ['Foo', 'a_function'], False, 3, id="function and class no __all__"
+						),
+				pytest.param(
+						testing_source_e, ['Foo', 'a_function'], True, 3, id="function and class with __all__"
+						),
+				pytest.param(
+						testing_source_f, ['Foo', 'a_function'],
+						True,
+						3,
+						id="function and class with __all__ and extra variable"
+						),
+				pytest.param(testing_source_g, ["a_function"], False, 3, id="async function no __all__"),
+				pytest.param(testing_source_h, [], False, 1, id="from import"),
+				pytest.param(testing_source_i, ["a_function"], False, 3, id="lots of lines"),
+				pytest.param(testing_source_j, ["a_function"], False, 14, id="multiline import"),
+				]
+		)
+def test_visitor_endlineno(source, members, found_all, last_import):
+	visitor = Visitor(True)
+	tree = ast.parse(source)
+	mark_text_ranges(tree, source)
+	visitor.visit(tree)
+
+	assert visitor.members == members
+	assert visitor.found_all is found_all
+	assert visitor.last_import is last_import
+
+
 @pytest.mark.parametrize(
 		"source, members, ret",
 		[
@@ -105,7 +146,6 @@ def test_visitor(source, members, found_all, last_import):
 				]
 		)
 def test_check_and_add_all(tmpdir, source, members: List[str], ret):
-
 	tmpfile = PathPlus(tmpdir) / "source.py"
 	tmpfile.write_text(source)
 
@@ -116,6 +156,7 @@ def test_check_and_add_all(tmpdir, source, members: List[str], ret):
 		assert f"__all__ = [{members_string}]" in tmpfile.read_text()
 
 
+@pytest.mark.skipif(condition=not (sys.version_info < (3, 8)), reason="Not required after python 3.8")
 @pytest.mark.parametrize(
 		"source, members, ret",
 		[
@@ -131,7 +172,6 @@ def test_check_and_add_all(tmpdir, source, members: List[str], ret):
 				]
 		)
 def test_check_and_add_all_single_quotes(tmpdir, source, members: List[str], ret):
-
 	tmpfile = PathPlus(tmpdir) / "source.py"
 	tmpfile.write_text(source)
 
@@ -146,7 +186,6 @@ def test_check_and_add_all_single_quotes(tmpdir, source, members: List[str], ret
 		pytest.param(mangled_source, [], id="mangled"),
 		])
 def test_check_and_add_all_mangled(tmpdir, capsys, source, members):
-
 	tmpfile = PathPlus(tmpdir) / "source.py"
 	tmpfile.write_text(source)
 	assert check_and_add_all(tmpfile) == 4

tests/test_utils.py

@@ -2,8 +2,7 @@
 import ast
 
 # this package
-from flake8_dunder_all import get_docstring_lineno
-from flake8_dunder_all.__main__ import tidy_docstring
+from flake8_dunder_all.utils import get_docstring_lineno, tidy_docstring
 
 
 def test_get_docstring_lineno():