Merge pull request #15 from PashaWNN/BANG_225__FIX_MULTIPLE_ASSIGNMENT_VALIDATION

BANG-225: fix validation of multiple assignments

Author: Arondit

.gitignore

@@ -102,3 +102,5 @@ venv.bak/
 
 # mypy
 .mypy_cache/
+
+.idea

flake8_variables_names/ast_helpers.py

@@ -1,19 +1,46 @@
 import ast
+import functools
 from typing import List, Tuple, Union
 
 from flake8_variables_names.list_helpers import flat
 
 
+@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]]:
-    if isinstance(assignment_node, ast.AnnAssign):
-        if isinstance(assignment_node.target, ast.Name):
-            return [(assignment_node.target.id, assignment_node.target)]
-    elif isinstance(assignment_node, ast.Assign):
-        names = [t for t in assignment_node.targets if isinstance(t, ast.Name)]
-        return [(n.id, n) for n in names]
-    return []
+    return [(n.id, n) for n in extract_names_from_node(assignment_node)]
 
 
 def get_var_names_from_funcdef(funcdef_node: ast.FunctionDef) -> List[Tuple[str, ast.arg]]:
@@ -25,22 +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)]
+    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])
-    ann_assignments = [n for n in ast.walk(ast_tree) if isinstance(n, ast.AnnAssign)]
-    var_info += flat([get_var_names_from_assignment(a) for a in ann_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_files/ok_names.py

@@ -1,2 +1,5 @@
 user_id = 23
 some_email = '[email protected]'
+some_normal_name, _ = -1, -2
+some_other_name, *_ = -1, -2, -3
+some_different_name, *other_different_name = -1, -2, -3

tests/test_files/short_names.py

@@ -2,3 +2,5 @@
 b = 1
 c = 3
 i = 4
+x, y, z = 5, 6, 7
+w, *q = 1, 2, 3

tests/test_variables_names.py

@@ -13,9 +13,20 @@ def test_ok_good_names():
 
 def test_ok_for_short_names_file():
     errors = run_validator_for_test_file('short_names.py', use_strict_mode=True)
-    assert len(errors) == 4
+    assert len(errors) == 9
     errors = run_validator_for_test_file('short_names.py', use_strict_mode=False)
+    assert len(errors) == 8
+    assert (
+        get_error_message(errors[0])
+        == "VNE001 single letter variable names like 'a' are not allowed"
+    )
+
+
+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"