infer-from-init

Author: Darshan808

IPython/core/guarded_eval.py

@@ -677,6 +677,10 @@ def dummy_function(*args, **kwargs):
         class_context = context.replace(transient_locals=class_locals)
         for child_node in node.body:
             eval_node(child_node, class_context)
+        # extract self.attribute assignments
+        init_attributes = _extract_init_attributes(node, class_context)
+        # Merge init attributes into class_locals
+        class_locals.update(init_attributes)
         bases = tuple([eval_node(base, context) for base in node.bases])
         dummy_class = type(node.name, bases, class_locals)
         context.transient_locals[node.name] = dummy_class
@@ -853,6 +857,116 @@ def dummy_function(*args, **kwargs):
     return None
 
 
+def _extract_init_attributes(class_node: ast.ClassDef, context: EvaluationContext):
+    """Extract attribute assignments from __init__ method.
+
+    Looks for patterns like:
+        self.attr = value
+        self.attr: Type = value
+
+    And infers their types by evaluating the assigned values.
+
+    Returns dictionary mapping attribute names to their inferred Duck types
+    """
+    attributes = {}
+    init_method = None
+    for node in class_node.body:
+        if isinstance(node, ast.FunctionDef) and node.name == "__init__":
+            init_method = node
+            break
+
+    if not init_method:
+        return attributes
+
+    temp_context = EvaluationContext(
+        globals=context.globals,
+        locals=context.locals,
+        evaluation=context.evaluation,
+        in_subscript=context.in_subscript,
+        transient_locals={},
+    )
+
+    for stmt in init_method.body:
+        # Handle regular assignments: self.attr = value
+        if isinstance(stmt, ast.Assign):
+            for target in stmt.targets:
+                if isinstance(target, ast.Attribute):
+                    if isinstance(target.value, ast.Name) and target.value.id == "self":
+                        attr_name = target.attr
+                        try:
+                            # Evaluate the assigned value
+                            value = eval_node(stmt.value, temp_context)
+                            if value is not None and value is not NOT_EVALUATED:
+                                inferred_type = _create_duck_from_value(value)
+                                if inferred_type is not None:
+                                    attributes[attr_name] = inferred_type
+                        except Exception:
+                            # Skip the attribute
+                            pass
+
+        # Handle annotated assignments: self.attr: Type = value
+        elif isinstance(stmt, ast.AnnAssign):
+            if isinstance(stmt.target, ast.Attribute):
+                if (
+                    isinstance(stmt.target.value, ast.Name)
+                    and stmt.target.value.id == "self"
+                ):
+                    attr_name = stmt.target.attr
+
+                    # Try to use the annotation
+                    if stmt.annotation:
+                        try:
+                            annotation = eval_node(stmt.annotation, temp_context)
+                            resolved = _resolve_annotation(annotation, context)
+                            if resolved is not None:
+                                attributes[attr_name] = resolved
+                                continue
+                        except Exception:
+                            pass
+
+                    # Try to infer from value
+                    if stmt.value:
+                        try:
+                            value = eval_node(stmt.value, temp_context)
+                            if value is not None and value is not NOT_EVALUATED:
+                                inferred_type = _create_duck_from_value(value)
+                                if inferred_type is not None:
+                                    attributes[attr_name] = inferred_type
+                        except Exception:
+                            pass
+
+    return attributes
+
+
+def _create_duck_from_value(value):
+    """Create a Duck object from an actual runtime value."""
+    if value is None or value is NOT_EVALUATED:
+        return None
+    value_type = type(value)
+    if isinstance(value, dict):
+        return _Duck(
+            attributes=dict.fromkeys(dir(dict())), items=value if value else {}
+        )
+    elif isinstance(value, list):
+        element_duck = None
+        if value:
+            element_duck = _create_duck_from_value(value[0])
+        return _Duck(
+            attributes=dict.fromkeys(dir(list())),
+            items=_GetItemDuck(lambda: element_duck),
+        )
+    elif isinstance(value, set):
+        return _Duck(attributes=dict.fromkeys(dir(set())))
+    elif isinstance(value, tuple):
+        return value
+    elif isinstance(value, (str, int, float, bool, bytes)):
+        return _Duck(attributes=dict.fromkeys(dir(value_type())))
+    else:
+        try:
+            return _create_duck_for_heap_type(value_type)
+        except Exception:
+            return _Duck(attributes=dict.fromkeys(dir(value)))
+
 def _eval_return_type(func: Callable, node: ast.Call, context: EvaluationContext):
     """Evaluate return type of a given callable function.