BANG-225: add more tests; refactor the code

Author: pshishmarev-bestdoctor

flake8_variables_names/ast_helpers.py

@@ -1,31 +1,44 @@
 import ast
+import functools
 from typing import List, Tuple, Union
 
 from flake8_variables_names.list_helpers import flat
 
 
-def extract_names_from_node(node: Union[ast.expr, ast.stmt]) -> List[ast.Name]:
-    if isinstance(node, ast.Name):
-        return [node]
-    if isinstance(node, ast.Assign):
-        nodes = []
-        for target in node.targets:
-            nodes.extend(extract_names_from_node(target))
-        return nodes
-    if isinstance(node, ast.AnnAssign):
-        return extract_names_from_node(node.target)
-    if isinstance(node, ast.Starred):
-        return extract_names_from_node(node.value)
-    if isinstance(node, ast.Tuple):
-        nodes = []
-        for elt in node.elts:
-            nodes.extend(extract_names_from_node(elt))
-        return nodes
+@functools.singledispatch
+def extract_names_from_node(node) -> List[ast.Name]:
     return []
 
 
+@extract_names_from_node.register
+def _extract_names_from_name_node(node: ast.Name):
+    return [node]
+
+
+@extract_names_from_node.register
+def _extract_names_from_assign_node(node: ast.Assign):
+    return flat([extract_names_from_node(target) for target in node.targets])
+
+
+# in some versions of Python, singledispatch does not support `Union` in type annotations
+@extract_names_from_node.register(ast.AnnAssign)
+@extract_names_from_node.register(ast.For)
+def _extract_names_from_annassign_node(node):
+    return extract_names_from_node(node.target)
+
+
+@extract_names_from_node.register
+def _extract_names_from_starred_node(node: ast.Starred):
+    return extract_names_from_node(node.value)
+
+
+@extract_names_from_node.register
+def _extract_names_from_tuple_node(node: ast.Tuple):
+    return flat([extract_names_from_node(elt) for elt in node.elts])
+
+
 def get_var_names_from_assignment(
-    assignment_node: Union[ast.Assign, ast.AnnAssign],
+    assignment_node: Union[ast.Assign, ast.AnnAssign, ast.For],
 ) -> List[Tuple[str, ast.AST]]:
     return [(n.id, n) for n in extract_names_from_node(assignment_node)]
 
@@ -39,20 +52,13 @@ def get_var_names_from_funcdef(funcdef_node: ast.FunctionDef) -> List[Tuple[str,
     return vars_info
 
 
-def get_var_names_from_for(for_node: ast.For) -> List[Tuple[str, ast.AST]]:
-    if isinstance(for_node.target, ast.Name):
-        return [(for_node.target.id, for_node.target)]
-    elif isinstance(for_node.target, ast.Tuple):
-        return [(n.id, n) for n in for_node.target.elts if isinstance(n, ast.Name)]
-    return []
-
-
 def extract_all_variable_names(ast_tree: ast.AST) -> List[Tuple[str, ast.AST]]:
     var_info: List[Tuple[str, ast.AST]] = []
-    assignments = [n for n in ast.walk(ast_tree) if isinstance(n, (ast.Assign, ast.AnnAssign))]
+    assignments = [
+        n for n in ast.walk(ast_tree)
+        if isinstance(n, (ast.Assign, ast.AnnAssign, ast.For))
+    ]
     var_info += flat([get_var_names_from_assignment(a) for a in assignments])
     funcdefs = [n for n in ast.walk(ast_tree) if isinstance(n, ast.FunctionDef)]
     var_info += flat([get_var_names_from_funcdef(f) for f in funcdefs])
-    fors = [n for n in ast.walk(ast_tree) if isinstance(n, ast.For)]
-    var_info += flat([get_var_names_from_for(f) for f in fors])
     return var_info

tests/test_files/loops_names.py

@@ -0,0 +1,5 @@
+for result in get_some_data():
+    pass
+
+for a, *b in zip(data, goosebumps):
+    print(a + b)

tests/test_variables_names.py

@@ -22,6 +22,17 @@ def test_ok_for_short_names_file():
     )
 
 
+def test_ok_for_names_in_loops_file():
+    errors = run_validator_for_test_file('loops_names.py', use_strict_mode=True)
+    assert len(errors) == 3
+    errors = run_validator_for_test_file('loops_names.py', use_strict_mode=False)
+    assert len(errors) == 2
+    assert (
+        get_error_message(errors[0])
+        == "VNE001 single letter variable names like 'a' are not allowed"
+    )
+
+
 def test_ok_for_commented_names_file():
     errors = run_validator_for_test_file('commented_names.py', use_strict_mode=True)
     assert not errors