analyses: type-guess when assigned from range

edulint/linting/analyses/types.py

@@ -6,7 +6,7 @@
 from astroid import nodes
 
 from edulint.linting.analyses.data_dependency import get_assigned_expressions
-from edulint.linting.checkers.utils import is_builtin, BaseVisitor
+from edulint.linting.checkers.utils import is_builtin, BaseVisitor, get_range_params
 
 
 class Type(Enum):
@@ -246,12 +246,16 @@ def visit_const(self, node: nodes.Const):
     def visit_name(self, node: nodes.Name):
         result = Types.empty()
         for assigned in get_assigned_expressions(
-            node, include_nodes=[], include_destructuring=False
+            node, include_nodes=[nodes.For, nodes.Comprehension], include_destructuring=True
         ):
             if assigned in self.visited_exprs:
                 continue
             self.visited_exprs.add(assigned)
-            result |= self.visit(assigned)
+            if isinstance(assigned.parent, (nodes.For, nodes.Comprehension)):
+                if get_range_params(assigned) is not None:
+                    result |= Types.int()
+            else:
+                result |= self.visit(assigned)
         return result
 
     def visit_import(self, _node: nodes.Import):

tests/test_local_defects.py

@@ -99,6 +99,11 @@ def test_local_custom(lines: List[str], expected_output: List[Problem]) -> None:
         "    b = ''",
         "    b = a + b",
     ], []),
+    ([
+        "def foo(a):",
+        "    for i in range(a):",
+        "        a = a + i",
+    ], [lazy_problem()]),
 ])
 def test_augassign_custom(lines: List[str], expected_output: List[Problem]) -> None:
     create_apply_and_lint(