Introduction of an interface for memory profile.

This patch touches #263.

Author: funilrys

PyFunceble/utils/profile.py

@@ -54,7 +54,9 @@
 import contextlib
 import cProfile
 import io
+import linecache
 import pstats
+import tracemalloc
 
 
 @contextlib.contextmanager
@@ -89,3 +91,43 @@ def profile_it(*, sort_stats: str = "cumulative", show_callers: bool = False):
         profiler_starts.print_callees()
 
     print(our_stream.getvalue())
+
+
+@contextlib.contextmanager
+def profile_memory(stats_mode: str = "lineno", top_limit: int = 10):
+    """
+    Provides a context manager which will activates memory profiling of our
+    source code.
+    """
+
+    tracemalloc.start()
+
+    yield
+
+    snapshot = tracemalloc.take_snapshot()
+    snapshot = snapshot.filter_traces(
+        (
+            tracemalloc.Filter(False, "<frozen importlib._bootstrap>"),
+            tracemalloc.Filter(False, "<unknown>"),
+            tracemalloc.Filter(False, "<frozen importlib._bootstrap_external>"),
+        )
+    )
+    top_stats = snapshot.statistics(stats_mode)
+
+    print("Top %s lines" % top_limit)
+    for index, stat in enumerate(top_stats[:top_limit], 1):
+        frame = stat.traceback[0]
+        print(
+            "#%s: %s:%s: %.1f KiB"
+            % (index, frame.filename, frame.lineno, stat.size / 1024)
+        )
+        line = linecache.getline(frame.filename, frame.lineno).strip()
+        if line:
+            print("    %s" % line)
+
+    other = top_stats[top_limit:]
+    if other:
+        size = sum(stat.size for stat in other)
+        print("%s other: %.1f KiB" % (len(other), size / 1024))
+    total = sum(stat.size for stat in top_stats)
+    print("Total allocated size: %.1f KiB" % (total / 1024))