From d5d7e08763dd1456d320ef5c482712e52575181a Mon Sep 17 00:00:00 2001
From: "John D. Pope" <jp@bellgeorge.com>
Date: Mon, 4 Nov 2024 13:49:19 +1100
Subject: [PATCH] Update memory_profiler.py

---
 memory_profiler.py | 139 +++++++++++++++++++++++++++++++++++----------
 1 file changed, 110 insertions(+), 29 deletions(-)

diff --git a/memory_profiler.py b/memory_profiler.py
index f0b52cd..4eede50 100644
--- a/memory_profiler.py
+++ b/memory_profiler.py
@@ -23,6 +23,12 @@
 import time
 import traceback
 import warnings
+from rich.console import Console
+from rich.table import Table
+from rich.syntax import Syntax
+from rich import box
+
+
 
 if sys.platform == "win32":
     # any value except signal.CTRL_C_EVENT and signal.CTRL_BREAK_EVENT
@@ -848,39 +854,112 @@ def enable(self):
 
     def disable(self):
         sys.settrace(self._original_trace_function)
-
-
 def show_results(prof, stream=None, precision=1):
-    if stream is None:
-        stream = sys.stdout
-    template = '{0:>6} {1:>12} {2:>12}  {3:>10}   {4:<}'
-
-    for (filename, lines) in prof.code_map.items():
-        header = template.format('Line #', 'Mem usage', 'Increment', 'Occurrences',
-                                 'Line Contents')
-
-        stream.write(u'Filename: ' + filename + '\n\n')
-        stream.write(header + u'\n')
-        stream.write(u'=' * len(header) + '\n')
+    """Display the profiling results using rich console for colored output."""
+    console = Console(file=stream) if stream else Console()
+    float_format = f'{{:.{precision}f}}'
 
+    for filename, line_iterator in prof.code_map.items():
+        console.print(f"\n[bold blue]Filename: {filename}[/bold blue]")
+        
+        # Create table
+        table = Table(
+            show_header=True,
+            header_style="bold white",
+            box=box.MINIMAL_DOUBLE_HEAD,
+            expand=True,
+            row_styles=["", "dim"]  # Alternating row styles
+        )
+        
+        # Add columns
+        table.add_column("Line #", justify="right", width=8)
+        table.add_column("Mem usage", justify="right", width=12)
+        table.add_column("Increment", justify="right", width=12)
+        table.add_column("Occurrences", justify="right", width=10)
+        table.add_column("Line Contents", ratio=1)
+
+        # Get all lines first
+        linecache.checkcache(filename)
         all_lines = linecache.getlines(filename)
-
-        float_format = u'{0}.{1}f'.format(precision + 4, precision)
-        template_mem = u'{0:' + float_format + '} MiB'
-        for (lineno, mem) in lines:
-            if mem:
-                inc = mem[0]
-                total_mem = mem[1]
-                total_mem = template_mem.format(total_mem)
-                occurrences = mem[2]
-                inc = template_mem.format(inc)
+        
+        # Convert line_iterator to list for multiple passes
+        try:
+            # First try direct iteration
+            lines_list = list(line_iterator)
+        except TypeError:
+            # If that fails, try calling items()
+            try:
+                lines_list = list(line_iterator.items())
+            except (AttributeError, TypeError):
+                console.print(f"[bold red]Error: Could not process lines for {filename}[/bold red]")
+                continue
+
+        # Get memory range for coloring
+        mem_values = []
+        for lineno, mem_info in lines_list:
+            if mem_info:
+                mem_values.append(mem_info[1])  # total memory
+        
+        max_mem = max(mem_values) if mem_values else 0
+
+        # Process each line
+        for lineno, mem_info in lines_list:
+            if not (0 < lineno <= len(all_lines)):
+                continue
+                
+            source_line = all_lines[lineno - 1].rstrip('\n')
+            
+            if mem_info:
+                increment, total_mem, occurrences = mem_info
+                
+                # Color based on memory usage
+                mem_ratio = total_mem / max_mem if max_mem else 0
+                color = ("green" if mem_ratio < 0.25 else
+                        "yellow" if mem_ratio < 0.5 else
+                        "orange1" if mem_ratio < 0.75 else
+                        "red")
+                
+                # Format the memory values with colors
+                mem_str = f"[{color}]{float_format.format(total_mem)}[/]"
+                inc_str = f"[{color}]{float_format.format(increment)}[/]"
+                
+                table.add_row(
+                    str(lineno),
+                    mem_str,
+                    inc_str,
+                    str(occurrences),
+                    Syntax(source_line, "python", theme="monokai", line_numbers=False)
+                )
             else:
-                total_mem = u''
-                inc = u''
-                occurrences = u''
-            tmp = template.format(lineno, total_mem, inc, occurrences, all_lines[lineno - 1])
-            stream.write(tmp)
-        stream.write(u'\n\n')
+                # Add row with no memory info
+                table.add_row(
+                    str(lineno),
+                    "", "", "",
+                    Syntax(source_line, "python", theme="monokai", line_numbers=False)
+                )
+
+        # Print the table
+        console.print(table)
+
+# Add debug function
+def print_code_map_structure(prof):
+    """Debug helper to print the structure of code_map"""
+    console = Console()
+    console.print("\n[bold red]Debug: Code Map Structure[/bold red]")
+    
+    for filename, line_iterator in prof.code_map.items():
+        console.print(f"\nFilename: {filename}")
+        console.print(f"Type of line_iterator: {type(line_iterator)}")
+        
+        # Try to get first few items
+        try:
+            if hasattr(line_iterator, 'items'):
+                first_items = list(itertools.islice(line_iterator.items(), 3))
+            else:
+                first_items = list(itertools.islice(line_iterator, 3))
+            console.print(f"First few items: {first_items}")
+        except Exception as e:
+            console.print(f"Error inspecting items: {str(e)}")
 
 
 def _func_exec(stmt, ns):
@@ -1160,6 +1239,8 @@ def load_ipython_extension(ip):
     MemoryProfilerMagics.register_magics(ip)
 
 
+
+
 def profile(func=None, stream=None, precision=1, backend='psutil'):
     """
     Decorator that will run the function and print a line-by-line profile