diff --git a/src/outline_tool.py b/src/outline_tool.py
index 9927f08..a533fde 100644
--- a/src/outline_tool.py
+++ b/src/outline_tool.py
@@ -53,7 +53,11 @@ class CodeOutliner:
      return doc.splitlines()[0]
    return None
 
+  count = [0]
   def walk(node: ast.AST, indent: int = 0) -> None:
+   count[0] += 1
+   if count[0] > 100000:
+    raise Exception("Infinite loop detected! " + str(type(node)))
    """
    [C: src/summarize.py:_summarise_python]
    """
@@ -70,15 +74,42 @@ class CodeOutliner:
     start_line = node.lineno
     end_line = getattr(node, "end_lineno", start_line)
     prefix = "[Async Func]" if isinstance(node, ast.AsyncFunctionDef) else "[Func]"
-    # Check if it's a method
-    # We can check the indent or the parent, but in AST walk we know if we are inside a ClassDef
-    # Let's use a simpler heuristic for the outline: if indent > 0, it's likely a method.
     if indent > 0:
      prefix = "[Method]"
-    output.append(f"{'  ' * indent}{prefix} {node.name} (Lines {start_line}-{end_line})")
+    returns = ""
+    if getattr(node, "returns", None):
+     try:
+      returns = f" -> {ast.unparse(node.returns)}"
+     except Exception:
+      pass
+    output.append(f"{'  ' * indent}{prefix} {node.name}{returns} (Lines {start_line}-{end_line})")
     doc = get_docstring(node)
     if doc:
      output.append(f"{'  ' * (indent + 1)}\"\"\"{doc}\"\"\"")
+    for item in node.body:
+     walk(item, indent + 1)
+   elif isinstance(node, ast.With):
+    is_imgui = False
+    try:
+     for item in node.items:
+      ctx_str = ast.unparse(item.context_expr)
+      if "imscope." in ctx_str or "imgui." in ctx_str:
+       start_line = node.lineno
+       end_line = getattr(node, "end_lineno", start_line)
+       output.append(f"{'  ' * indent}[ImGui Scope] {ctx_str} (Lines {start_line}-{end_line})")
+       is_imgui = True
+       break
+    except Exception:
+     pass
+    for item in node.body:
+     walk(item, indent + 1 if is_imgui else indent)
+   else:
+    for block_attr in ("body", "orelse", "handlers", "finalbody"):
+     block = getattr(node, block_attr, [])
+     if isinstance(block, list):
+      for item in block:
+       walk(item, indent)
+  
   for node in tree.body:
    walk(node)
   return "\n".join(output)
diff --git a/tests/test_outline_tool.py b/tests/test_outline_tool.py
new file mode 100644
index 0000000..034ff0b
--- /dev/null
+++ b/tests/test_outline_tool.py
@@ -0,0 +1,42 @@
+import pytest
+import ast
+from pathlib import Path
+from src.outline_tool import CodeOutliner
+
+def test_code_outliner_type_hints():
+    code = """
+class Test:
+    def my_method(self, x: int) -> str:
+        return "hello"
+        
+def my_func(a: bool) -> None:
+    pass
+"""
+    outliner = CodeOutliner()
+    result = outliner.outline(code)
+    assert "[Method] my_method -> str (Lines 3-4)" in result
+    assert "[Func] my_func -> None (Lines 6-7)" in result
+
+def test_code_outliner_imgui_scopes():
+    code = """
+def render_gui() -> None:
+    with imscope.window("My Window"):
+        imgui.text("Hello")
+        with imscope.child("My Child"):
+            pass
+"""
+    outliner = CodeOutliner()
+    result = outliner.outline(code)
+    assert "[ImGui Scope] imscope.window('My Window') (Lines 3-6)" in result
+    assert "[ImGui Scope] imscope.child('My Child') (Lines 5-6)" in result
+
+def test_code_outliner_nested_ifs():
+    code = """
+def render_gui() -> None:
+    if True:
+        with imscope.window("My Window"):
+            pass
+"""
+    outliner = CodeOutliner()
+    result = outliner.outline(code)
+    assert "[ImGui Scope] imscope.window('My Window') (Lines 4-5)" in result