BANG-225: add more corner cases

Author: pshishmarev-bestdoctor

flake8_variables_names/ast_helpers.py

@@ -4,27 +4,30 @@
 from flake8_variables_names.list_helpers import flat
 
 
-def extract_name_nodes_from_assignment(assignment_node: ast.Assign) -> List[ast.Name]:
-    names = []
-    for target in assignment_node.targets:
-        if isinstance(target, ast.Name):
-            names.append(target)
-        elif isinstance(target, ast.Tuple):
-            names.extend([dim for dim in target.dims if isinstance(dim, ast.Name)])
-    return names
+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
+    return []
 
 
 def get_var_names_from_assignment(
     assignment_node: Union[ast.Assign, ast.AnnAssign],
 ) -> List[Tuple[str, ast.AST]]:
-    if (
-        isinstance(assignment_node, ast.AnnAssign)
-        and isinstance(assignment_node.target, ast.Name)
-    ):
-        return [(assignment_node.target.id, assignment_node.target)]
-    elif isinstance(assignment_node, ast.Assign):
-        return [(n.id, n) for n in extract_name_nodes_from_assignment(assignment_node)]
-    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]]:
@@ -46,10 +49,8 @@ def get_var_names_from_for(for_node: ast.For) -> List[Tuple[str, ast.AST]]:
 
 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))]
     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)]

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

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

tests/test_variables_names.py

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