Correctly handle @overload decorated functions.

Author: domdfcoding

flake8_dunder_all/__init__.py

@@ -32,7 +32,7 @@
 # stdlib
 import ast
 import sys
-from typing import Any, Generator, List, Tuple, Type, Union
+from typing import Any, Generator, List, Set, Tuple, Type, Union
 
 # 3rd party
 from consolekit.terminal_colours import Fore
@@ -64,11 +64,12 @@ class Visitor(ast.NodeVisitor):
 
 	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
+	members: Set[str]  #: List of functions and classed defined in the AST
+	use_endlineno: bool
 
 	def __init__(self, use_endlineno: bool = False) -> None:
 		self.found_all = False
-		self.members: List[str] = []
+		self.members = set()
 		self.last_import = 0
 		self.use_endlineno = use_endlineno
 
@@ -91,8 +92,30 @@ def handle_def(self, node: Union[ast.FunctionDef, ast.AsyncFunctionDef, ast.Clas
 		:param node: The node being visited.
 		"""
 
-		if not node.name.startswith('_'):
-			self.members.append(node.name)
+		decorators = []
+
+		for deco in node.decorator_list:
+			if isinstance(deco, ast.Name):
+				decorators.append(deco.id)
+			elif isinstance(deco, ast.Attribute):
+				parts = [deco.attr]
+
+				# last_part = deco.value
+				#
+				# while True:
+				# 	if isinstance(last_part, ast.Attribute):
+				# 		parts.append(last_part.attr)
+				# 		last_part = last_part.value
+				# 	elif isinstance(last_part, ast.Name):
+				# 		parts.append(last_part.id)
+				# 		break
+				# 	else:
+				# 		break
+
+				decorators.append('.'.join(reversed(parts)))
+
+		if not node.name.startswith('_') and "overload" not in decorators:
+			self.members.add(node.name)
 
 	def visit_FunctionDef(self, node: ast.FunctionDef):
 		"""
@@ -236,7 +259,7 @@ def check_and_add_all(filename: PathPlus, quote_type: str = '"') -> int:
 		if not visitor.members:
 			return 0
 
-		members = repr(visitor.members).replace(bad_quote, quote_type)
+		members = repr(sorted(visitor.members)).replace(bad_quote, quote_type)
 
 		lines = filename.read_text().split('\n')
 

tests/common.py

@@ -100,6 +100,20 @@ async def a_function(): ...
 async def a_function(): ...
 """
 
+testing_source_k = '''
+"""a docstring"""
+
+@overload
+def a_function(): ...
+'''
+
+testing_source_l = '''
+"""a docstring"""
+
+@typing.overload
+def a_function(): ...
+'''
+
 mangled_source = '''
 """a docstring
 import foo

tests/test_flake8_dunder_all.py

@@ -24,7 +24,9 @@
 		testing_source_g,
 		testing_source_h,
 		testing_source_i,
-		testing_source_j
+		testing_source_j,
+		testing_source_k,
+		testing_source_l
 		)
 
 
@@ -84,7 +86,7 @@ def test_visitor(source, members, found_all, last_import):
 	visitor = Visitor()
 	visitor.visit(ast.parse(source))
 
-	assert visitor.members == members
+	assert sorted(visitor.members) == members
 	assert visitor.found_all is found_all
 	assert visitor.last_import is last_import
 
@@ -121,7 +123,7 @@ def test_visitor_endlineno(source, members, found_all, last_import):
 	mark_text_ranges(tree, source)
 	visitor.visit(tree)
 
-	assert visitor.members == members
+	assert sorted(visitor.members) == members
 	assert visitor.found_all is found_all
 	assert visitor.last_import is last_import
 
@@ -144,6 +146,8 @@ def test_visitor_endlineno(source, members, found_all, last_import):
 				pytest.param(testing_source_g, ["a_function"], 1, id="async function no __all__"),
 				pytest.param(testing_source_h, [], 0, id="from import"),
 				pytest.param(testing_source_i, [], 1, id="lots of lines"),
+				pytest.param(testing_source_k, [], 0, id="overload"),
+				pytest.param(testing_source_l, [], 0, id="typing.overload"),
 				]
 		)
 def test_check_and_add_all(tmpdir, source, members: List[str], ret):