1733 add lambda class assignment violation

tests/test_visitors/test_ast/test_classes/test_instance_assignment.py

@@ -4,42 +4,62 @@
     InstanceLambdaAssignmentViolation,
 )
 from wemake_python_styleguide.visitors.ast.classes import (
-    InstanceAssignmentVisitor,
+    AttributesAssignmentVisitor,
 )
 
 instance_lambda_assignment = """
 class Example(object):
     def __init__(self):
-        {0}.{1} = lambda: ...
+        {0}
 """
 
+class_lambda_assignment = """
+class Example(object):
+    {0}
+"""
 
-@pytest.mark.parametrize('reference', [
-    'self',
-    'cls',
-    'mcs',
-    'other',
-])
-@pytest.mark.parametrize('attr', [
-    'attr',
-    '_attr',
-    '__attr',
+
+@pytest.mark.parametrize('assignment', [
+    'self.attr = lambda: ...',
+    'self.attr1 = self.attr2 = lambda: ...',
+    'self.attr1, self.attr2 = lambda: ..., "value"',
+    'self.attr: Callable[[], None] = lambda: ...',
 ])
 def test_instance_lambda_assignment(
     assert_errors,
     assert_error_text,
     parse_ast_tree,
     default_options,
-    reference,
-    attr,
+    assignment,
     mode,
 ):
     """Testing lambda assignment to instance."""
-    tree = parse_ast_tree(mode(instance_lambda_assignment.format(
-        reference, attr,
-    )))
+    tree = parse_ast_tree(mode(instance_lambda_assignment.format(assignment)))
+
+    visitor = AttributesAssignmentVisitor(default_options, tree=tree)
+    visitor.run()
+
+    assert_errors(visitor, [InstanceLambdaAssignmentViolation])
+
+
+@pytest.mark.parametrize('assignment', [
+    'attr = lambda: ...',
+    'attr1 = attr2 = lambda: ...',
+    'attr1, attr2 = lambda: ..., "value"',
+    'attr: Callable[[], None] = lambda: ...',
+])
+def test_class_lambda_assignment(
+    assert_errors,
+    assert_error_text,
+    parse_ast_tree,
+    default_options,
+    assignment,
+    mode,
+):
+    """Testing lambda assignment to class."""
+    tree = parse_ast_tree(mode(class_lambda_assignment.format(assignment)))
 
-    visitor = InstanceAssignmentVisitor(default_options, tree=tree)
+    visitor = AttributesAssignmentVisitor(default_options, tree=tree)
     visitor.run()
 
     assert_errors(visitor, [InstanceLambdaAssignmentViolation])

wemake_python_styleguide/logic/tree/assignments.py

@@ -0,0 +1,87 @@
+import ast
+import itertools
+from typing import Iterable, List, Optional, Tuple
+
+from wemake_python_styleguide import constants
+from wemake_python_styleguide.compat.aliases import AssignNodes
+from wemake_python_styleguide.compat.functions import get_assign_targets
+from wemake_python_styleguide.logic import nodes
+from wemake_python_styleguide.types import AnyAssign
+
+#: Type alias for the assignments we return from class inspection.
+_AllAssignments = Tuple[List[AnyAssign], List[AnyAssign]]
+
+
+def flat_assignment_values(assigns: Iterable[AnyAssign]) -> Iterable[ast.AST]:
+    """
+    Returns flat values from assignment.
+
+    Use this function when you need to get list of values
+    from assign nodes.
+    """
+    return itertools.chain.from_iterable((
+        _flat_nodes(assign.value)
+        for assign in assigns
+        if isinstance(assign.value, ast.AST)
+    ))
+
+
+def _flat_nodes(node: ast.AST) -> List[ast.AST]:
+    flatten_nodes: List[ast.AST] = []
+
+    if isinstance(node, ast.Tuple):
+        for subnode in node.elts:
+            flatten_nodes.extend(_flat_nodes(subnode))
+    else:
+        flatten_nodes.append(node)
+    return flatten_nodes
+
+
+def get_assignments(node: ast.ClassDef) -> _AllAssignments:
+    """
+    Helper to get all assignments from class nod definitions.
+
+    Args:
+        node: class node definition.
+
+    Returns:
+        A tuple of lists for both class and instance level variables.
+    """
+    class_assignments = []
+    instance_assignments = []
+
+    for subnode in ast.walk(node):
+        instance_assign = _get_instance_assignment(subnode)
+        if instance_assign is not None:
+            instance_assignments.append(instance_assign)
+            continue
+
+        class_assign = _get_class_assignment(node, subnode)
+        if class_assign is not None:
+            class_assignments.append(class_assign)
+
+    return class_assignments, instance_assignments
+
+
+def _get_instance_assignment(subnode: ast.AST) -> Optional[AnyAssign]:
+    return subnode if (
+        isinstance(subnode, AssignNodes) and
+        any(
+            isinstance(target, ast.Attribute) and
+            isinstance(target.value, ast.Name) and
+            target.value.id in constants.SPECIAL_ARGUMENT_NAMES_WHITELIST
+            for targets in get_assign_targets(subnode)
+            for target in _flat_nodes(targets)
+        )
+    ) else None
+
+
+def _get_class_assignment(
+    node: ast.ClassDef,
+    subnode: ast.AST,
+) -> Optional[AnyAssign]:
+    return subnode if (
+        isinstance(subnode, AssignNodes) and
+        nodes.get_context(subnode) is node and
+        getattr(subnode, 'value', None)
+    ) else None

wemake_python_styleguide/presets/types/tree.py

@@ -93,7 +93,7 @@
     classes.WrongSlotsVisitor,
     classes.ClassAttributeVisitor,
     classes.ClassMethodOrderVisitor,
-    classes.InstanceAssignmentVisitor,
+    classes.AttributesAssignmentVisitor,
 
     blocks.BlockVariableVisitor,
     blocks.AfterBlockVariablesVisitor,

wemake_python_styleguide/visitors/ast/blocks.py

@@ -76,7 +76,7 @@ class BlockVariableVisitor(base.BaseNodeVisitor):
     )
 
     _scope_predicates: Tuple[_ScopePredicate, ...] = (
-        lambda node, names: predicates.is_property_setter(node),
+        lambda node, names: predicates.is_property_setter(node),  # noqa: WPS467
         predicates.is_same_value_reuse,
         predicates.is_same_try_except_cases,
     )

wemake_python_styleguide/visitors/ast/classes.py

@@ -1,4 +1,5 @@
 import ast
+import itertools
 from collections import defaultdict
 from typing import ClassVar, DefaultDict, FrozenSet, List, Optional
 
@@ -11,6 +12,7 @@
 from wemake_python_styleguide.logic.arguments import function_args, super_args
 from wemake_python_styleguide.logic.naming import access, name_nodes
 from wemake_python_styleguide.logic.tree import (
+    assignments,
     attributes,
     classes,
     functions,
@@ -435,12 +437,12 @@ def _ideal_order(self, first: str) -> int:
 
 
 @final
-class InstanceAssignmentVisitor(base.BaseNodeVisitor):
-    """Checks that all attributes of instance are assigned correctly."""
+class AttributesAssignmentVisitor(base.BaseNodeVisitor):
+    """Checks that all attributes inside class are assigned correctly."""
 
     def visit_ClassDef(self, node: ast.ClassDef) -> None:
         """
-        Ensures that instance has correct attributes assignment.
+        Ensures that class or instance has correct attributes assignment.
 
         Raises:
             InstanceLambdaAssignment
@@ -450,10 +452,17 @@ def visit_ClassDef(self, node: ast.ClassDef) -> None:
         self.generic_visit(node)
 
     def _check_lambda_assignment(self, node: ast.ClassDef) -> None:
-        for subnode in walk.get_subnodes_by_type(node, ast.Assign):
-            for target in subnode.targets:
-                is_lambda = isinstance(subnode.value, ast.Lambda)
-                if isinstance(target, ast.Attribute) and is_lambda:
-                    self.add_violation(
-                        bp.InstanceLambdaAssignmentViolation(subnode),
-                    )
+        class_assignments, instance_assignments = assignments.get_assignments(
+            node,
+        )
+        flatten_values = assignments.flat_assignment_values(
+            itertools.chain(
+                class_assignments,
+                instance_assignments,
+            ),
+        )
+        for assigned_value in flatten_values:
+            if isinstance(assigned_value, ast.Lambda):
+                self.add_violation(
+                    bp.InstanceLambdaAssignmentViolation(assigned_value),
+                )