add-loops-and-conditions

Author: Darshan808

IPython/core/guarded_eval.py

@@ -817,6 +817,82 @@ def dummy_function(*args, **kwargs):
         if hasattr(value, "__awaited_type__"):
             return value.__awaited_type__
         return value
+    if isinstance(node, ast.While):
+        loop_locals = context.transient_locals.copy()
+        loop_context = context.replace(transient_locals=loop_locals)
+
+        result = None
+        for stmt in node.body:
+            result = eval_node(stmt, loop_context)
+
+        policy = get_policy(context)
+        merged_locals = _merge_dicts_by_key(
+            [loop_locals, context.transient_locals.copy()], policy
+        )
+        context.transient_locals.update(merged_locals)
+
+        return result
+    if isinstance(node, ast.For):
+        try:
+            iterable = eval_node(node.iter, context)
+        except:
+            iterable = None
+
+        loop_var_value = None
+        if iterable is not None:
+            try:
+                if policy.can_call(iterable.__iter__):
+                    iterator = iter(iterable)
+                    loop_var_value = next(iterator)
+            except:
+                pass
+
+        loop_locals = context.transient_locals.copy()
+        if isinstance(node.target, ast.Name):
+            if loop_var_value is not None:
+                loop_locals[node.target.id] = loop_var_value
+
+        loop_context = context.replace(transient_locals=loop_locals)
+
+        result = None
+        for stmt in node.body:
+            result = eval_node(stmt, loop_context)
+
+        policy = get_policy(context)
+        merged_locals = _merge_dicts_by_key(
+            [loop_locals, context.transient_locals.copy()], policy
+        )
+        context.transient_locals.update(merged_locals)
+
+        return result
+    if isinstance(node, ast.If):
+        branches = []
+        current = node
+        result = None
+        while True:
+            branch_locals = context.transient_locals.copy()
+            branch_context = context.replace(transient_locals=branch_locals)
+            for stmt in current.body:
+                result = eval_node(stmt, branch_context)
+            branches.append(branch_locals)
+            if not current.orelse:
+                break
+            elif len(current.orelse) == 1 and isinstance(current.orelse[0], ast.If):
+                # It's an elif - continue loop
+                current = current.orelse[0]
+            else:
+                # It's an else block - process and break
+                else_locals = context.transient_locals.copy()
+                else_context = context.replace(transient_locals=else_locals)
+                for stmt in current.orelse:
+                    result = eval_node(stmt, else_context)
+                branches.append(else_locals)
+                break
+        branches.append(context.transient_locals.copy())
+        policy = get_policy(context)
+        merged_locals = _merge_dicts_by_key(branches, policy)
+        context.transient_locals.update(merged_locals)
+        return result
     if isinstance(node, ast.Assign):
         return _handle_assign(node, context)
     if isinstance(node, ast.AnnAssign):
@@ -1006,6 +1082,23 @@ def dummy_function(*args, **kwargs):
     return None
 
 
+def _merge_dicts_by_key(dicts: list, policy: EvaluationPolicy):
+    """Merge multiple dictionaries, combining values for each key."""
+    if len(dicts) == 1:
+        return dicts[0]
+
+    all_keys = set()
+    for d in dicts:
+        all_keys.update(d.keys())
+
+    merged = {}
+    for key in all_keys:
+        values = [d[key] for d in dicts if key in d]
+        if values:
+            merged[key] = _merge_values(values, policy)
+
+    return merged
+
 def _merge_values(values, policy: EvaluationPolicy):
     """Recursively merge multiple values, combining attributes and dict items."""
     if len(values) == 1:
@@ -1368,6 +1461,7 @@ def _list_methods(cls, source=None):
     collections.Counter.most_common,
     object.__dir__,
     type.__dir__,
+    _Duck.__dir__,
 }
 
 BUILTIN_GETATTR: set[MayHaveGetattr] = {

tests/test_completer.py

@@ -2316,6 +2316,120 @@ def _(expected):
             ),
             "append",
         ],
+        [
+            "\n".join(["t = []", "if some_condition:", "    t."]),
+            "append",
+        ],
+        [
+            "\n".join(
+                [
+                    "t = []",
+                    "if some_condition:",
+                    "    t = 'string'",
+                    "t.",
+                ]
+            ),
+            ["append", "capitalize"],
+        ],
+        [
+            "\n".join(
+                [
+                    "t = []",
+                    "if some_condition:",
+                    "    t = 'string'",
+                    "else:",
+                    "    t.",
+                ]
+            ),
+            "append",
+        ],
+        [
+            "\n".join(
+                [
+                    "t = []",
+                    "if some_condition:",
+                    "    t = 'string'",
+                    "else:",
+                    "    t = 1",
+                    "t.",
+                ]
+            ),
+            ["append", "capitalize", "as_integer_ratio"],
+        ],
+        [
+            "\n".join(
+                [
+                    "t = []",
+                    "if condition_1:",
+                    "    t = 'string'",
+                    "elif condition_2:",
+                    "    t = 1",
+                    "elif condition_3:",
+                    "    t.",
+                ]
+            ),
+            "append",
+        ],
+        [
+            "\n".join(
+                [
+                    "t = []",
+                    "if condition_1:",
+                    "    t = 'string'",
+                    "elif condition_2:",
+                    "    t = 1",
+                    "elif condition_3:",
+                    "    t = {}",
+                    "t.",
+                ]
+            ),
+            ["append", "capitalize", "as_integer_ratio", "keys"],
+        ],
+        [
+            "\n".join(
+                [
+                    "t = []",
+                    "if condition_1:",
+                    "    if condition_2:",
+                    "        t = 'nested'",
+                    "t.",
+                ]
+            ),
+            ["append", "capitalize"],
+        ],
+        [
+            "\n".join(
+                [
+                    "a = []",
+                    "while condition:",
+                    "    a.",
+                ]
+            ),
+            "append",
+        ],
+        [
+            "\n".join(
+                [
+                    "t = []",
+                    "while condition:",
+                    "    t = 'str'",
+                    "t.",
+                ]
+            ),
+            ["append", "capitalize"],
+        ],
+        [
+            "\n".join(
+                [
+                    "t = []",
+                    "while condition_1:",
+                    "    while condition_2:",
+                    "        t = 'str'",
+                    "t.",
+                ]
+            ),
+            ["append", "capitalize"],
+        ],
     ],
 )
 def test_undefined_variables(use_jedi, evaluation, code, insert_text):