├── docs ├── index.md └── terminalui_showcase.png ├── lblprof ├── tests │ ├── example_modules │ │ ├── module_1 │ │ │ ├── __init__.py │ │ │ └── script1.py │ │ └── module_2 │ │ │ └── script2.py │ ├── example_scripts │ │ ├── simple_sleep.py │ │ ├── nested_functions.py │ │ ├── list_comprehension.py │ │ ├── double_zip.py │ │ ├── import_pandas.py │ │ ├── try_except.py │ │ ├── fibonacci.py │ │ ├── separate_frames.py │ │ ├── generator_function.py │ │ ├── complex_example.py │ │ ├── chroma_vector_search.py │ │ └── data_computation.py │ ├── test_tracing_no_new_context.py │ ├── demo_example_script.py │ ├── test_line_stats_tree.py │ └── test_tree_coherence.py ├── benchmark │ ├── funcs.py │ ├── bench_isolated_process.py │ └── benchmark_overhead.py ├── __init__.py ├── line_stat_object.py ├── custom_sysmon.py ├── curses_ui.py └── line_stats_tree.py ├── .gitignore ├── Makefile ├── .pre-commit-config.yaml ├── mkdocs.yml ├── .github └── workflows │ ├── python-checks.yml │ └── mkdocs.yml ├── pyproject.toml └── Readme.md /docs/index.md: -------------------------------------------------------------------------------- 1 | # Welcome to LblProf -------------------------------------------------------------------------------- /lblprof/tests/example_modules/module_1/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /docs/terminalui_showcase.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/le-codeur-rapide/lblprof/HEAD/docs/terminalui_showcase.png -------------------------------------------------------------------------------- /lblprof/tests/example_scripts/simple_sleep.py: -------------------------------------------------------------------------------- 1 | import time 2 | 3 | if __name__ == "__main__": 4 | time.sleep(0.2) 5 | -------------------------------------------------------------------------------- /lblprof/benchmark/funcs.py: -------------------------------------------------------------------------------- 1 | def import_pandas(): 2 | import pandas # noqa 3 | 4 | 5 | sample_functions = [import_pandas] 6 | -------------------------------------------------------------------------------- /lblprof/tests/example_modules/module_1/script1.py: -------------------------------------------------------------------------------- 1 | import time 2 | 3 | time.sleep(0.1) 4 | 5 | 6 | def function_1(): 7 | time.sleep(0.1) 8 | -------------------------------------------------------------------------------- /lblprof/tests/example_modules/module_2/script2.py: -------------------------------------------------------------------------------- 1 | import time 2 | 3 | time.sleep(0.1) 4 | 5 | 6 | def function_2(): 7 | time.sleep(0.1) 8 | -------------------------------------------------------------------------------- /.gitignore: -------------------------------------------------------------------------------- 1 | .vscode/ 2 | __pycache__/ 3 | lbl_profiler.egg-info/ 4 | .env 5 | dist/ 6 | build/ 7 | *.egg-info/ 8 | uv.lock 9 | *.csv 10 | *.json 11 | *.txt 12 | *.ipynb 13 | tests_examples/ 14 | site/ -------------------------------------------------------------------------------- /lblprof/tests/example_scripts/nested_functions.py: -------------------------------------------------------------------------------- 1 | import time 2 | 3 | 4 | def outer(): 5 | time.sleep(0.05) 6 | 7 | def inner(): 8 | time.sleep(0.05) 9 | 10 | inner() 11 | 12 | 13 | if __name__ == "__main__": 14 | outer() 15 | -------------------------------------------------------------------------------- /lblprof/tests/example_scripts/list_comprehension.py: -------------------------------------------------------------------------------- 1 | import time 2 | 3 | 4 | def main(): 5 | _ = [time.sleep(0.1) for _ in range(2)] 6 | 7 | _ = [time.sleep(0.1) for _ in range(2)] 8 | return 9 | 10 | 11 | if __name__ == "__main__": 12 | main() 13 | -------------------------------------------------------------------------------- /Makefile: -------------------------------------------------------------------------------- 1 | 2 | # Default target 3 | .PHONY: all 4 | all: lint test 5 | 6 | install: 7 | uv venv && uv pip install -e . 8 | 9 | test: 10 | uv run pytest 11 | 12 | lint: 13 | uv run ruff check . --fix 14 | 15 | benchmark: 16 | uv run lblprof/benchmark/benchmark_overhead.py -------------------------------------------------------------------------------- /lblprof/tests/example_scripts/double_zip.py: -------------------------------------------------------------------------------- 1 | def doublezip(): 2 | rows = [(1, 2, None), (3, 4, None), (5, 6, None), (7, 8, None)] 3 | _ = list(zip(*[col for col in zip(*rows) if any(cell is not None for cell in col)])) 4 | _ = any(cell is not None for cell in rows[0]) 5 | 6 | 7 | if __name__ == "__main__": 8 | doublezip() 9 | -------------------------------------------------------------------------------- /lblprof/tests/example_scripts/import_pandas.py: -------------------------------------------------------------------------------- 1 | import time 2 | 3 | 4 | def main(): 5 | start = time.perf_counter() 6 | import pandas as pd # noqa: F401 7 | 8 | end = time.perf_counter() 9 | print(f"Time taken to import pandas: {end - start} seconds") 10 | return 11 | 12 | 13 | if __name__ == "__main__": 14 | main() 15 | -------------------------------------------------------------------------------- /lblprof/tests/example_scripts/try_except.py: -------------------------------------------------------------------------------- 1 | import time 2 | 3 | 4 | def main(): 5 | time.sleep(0.1) 6 | try: 7 | time.sleep(0.1) 8 | raise Exception("test") 9 | except Exception: 10 | time.sleep(0.1) 11 | finally: 12 | time.sleep(0.1) 13 | time.sleep(0.1) 14 | return 15 | 16 | 17 | if __name__ == "__main__": 18 | main() 19 | -------------------------------------------------------------------------------- /lblprof/tests/example_scripts/fibonacci.py: -------------------------------------------------------------------------------- 1 | import time 2 | 3 | 4 | def fib_iterative(n): 5 | if n < 0: 6 | raise ValueError("Input must be a non-negative integer.") 7 | a, b = 0, 1 8 | for _ in range(n): 9 | time.sleep(0.05) 10 | a, b = b, a + b 11 | return a 12 | 13 | 14 | if __name__ == "__main__": 15 | time.sleep(0.05) 16 | fib_iterative(4) 17 | -------------------------------------------------------------------------------- /.pre-commit-config.yaml: -------------------------------------------------------------------------------- 1 | repos: 2 | - repo: https://github.com/psf/black 3 | rev: 24.2.0 # Use the latest stable version 4 | hooks: 5 | - id: black 6 | files: ^.*\.py$ 7 | 8 | - repo: https://github.com/astral-sh/ruff-pre-commit 9 | rev: v0.3.0 # Use the latest stable version 10 | hooks: 11 | - id: ruff 12 | args: [--fix] 13 | files: ^.*\.py$ 14 | -------------------------------------------------------------------------------- /mkdocs.yml: -------------------------------------------------------------------------------- 1 | site_name: LblProf 2 | site_description: Line by line terminal based profiler 3 | theme: 4 | name: mkdocs 5 | features: 6 | - navigation.tabs 7 | - navigation.sections 8 | - navigation.top 9 | - search.highlight 10 | - search.share 11 | 12 | nav: 13 | - Home: index.md 14 | - Documentation: documentation.md 15 | - Changelog: changelog.md 16 | 17 | markdown_extensions: 18 | - pymdownx.highlight 19 | - pymdownx.superfences -------------------------------------------------------------------------------- /lblprof/tests/test_tracing_no_new_context.py: -------------------------------------------------------------------------------- 1 | # See https://github.com/le-codeur-rapide/lblprof/issues/5 for context 2 | 3 | import time 4 | from lblprof import start_tracing, stop_tracing, tracer 5 | 6 | 7 | def test_tracing_no_new_context(): 8 | start_tracing() 9 | time.sleep(0.1) 10 | stop_tracing() 11 | # check that the time.sleep is in the tree 12 | line_codes = [line.source for line in tracer.tree.root_lines] 13 | assert "time.sleep(0.1)" in line_codes 14 | -------------------------------------------------------------------------------- /lblprof/tests/example_scripts/separate_frames.py: -------------------------------------------------------------------------------- 1 | import time 2 | 3 | # from lblprof import start_tracing, stop_tracing, show_tree 4 | import logging 5 | 6 | logging.basicConfig(level=logging.DEBUG) 7 | 8 | 9 | def main(): 10 | 11 | time.sleep(0.1) 12 | main2() 13 | 14 | 15 | def main2(): 16 | time.sleep(0.1) 17 | # stop_tracing() 18 | time.sleep(0.1) 19 | 20 | 21 | if __name__ == "__main__": 22 | # start_tracing() 23 | main() 24 | time.sleep(0.1) 25 | 26 | # show_tree() 27 | -------------------------------------------------------------------------------- /lblprof/tests/example_scripts/generator_function.py: -------------------------------------------------------------------------------- 1 | # import time 2 | # def generator(): 3 | # for i in range(10): 4 | # time.sleep(0.05) 5 | # yield i 6 | 7 | # if __name__ == "__main__": 8 | # for i in generator(): 9 | # print(i) 10 | 11 | # gene = list(generator()) 12 | 13 | # pure_gene = list(time.sleep(0.1) for i in range(10)) 14 | 15 | # double_gen = list( 16 | # time.sleep(0.1) for i in range(2) for j in range(2) if True 17 | # ) 18 | 19 | # # rows = [(1, 2, None), (3, 4, None), (5, 6, None)] 20 | # # filtered = list( 21 | # # zip(*[ 22 | # # col for col in zip(*rows) 23 | # # if any(cell is not None for cell in col) 24 | # # ]) 25 | # # ) 26 | -------------------------------------------------------------------------------- /.github/workflows/python-checks.yml: -------------------------------------------------------------------------------- 1 | name: Python Checks 2 | 3 | on: 4 | push: 5 | branches: [ main ] 6 | pull_request: 7 | branches: [ main ] 8 | 9 | jobs: 10 | test: 11 | runs-on: ubuntu-latest 12 | 13 | steps: 14 | - uses: actions/checkout@v4 15 | 16 | - name: Set up Python 17 | uses: actions/setup-python@v4 18 | with: 19 | python-version: '3.12' 20 | 21 | - name: Install dependencies 22 | run: | 23 | python -m pip install --upgrade pip 24 | pip install uv 25 | uv venv 26 | uv pip install .[dev] 27 | 28 | - name: Run Black 29 | run: uv run black . --check 30 | 31 | - name: Run Ruff 32 | run: uv run ruff check . 33 | 34 | - name: Run tests 35 | run: uv run pytest -------------------------------------------------------------------------------- /lblprof/benchmark/bench_isolated_process.py: -------------------------------------------------------------------------------- 1 | """This script is run in an isolated process to benchmark the overhead of lblprof. 2 | It takes as arguments the module name, function name, and mode (profiled or unprofiled)""" 3 | 4 | if __name__ == "__main__": 5 | import sys 6 | import time 7 | import json 8 | import importlib 9 | 10 | # args: module function mode 11 | module, func_name, mode = sys.argv[1], sys.argv[2], sys.argv[3] 12 | 13 | mod = importlib.import_module(module) 14 | fn = getattr(mod, func_name) 15 | 16 | start = time.perf_counter() 17 | if mode == "profiled": 18 | from lblprof import start_tracing, stop_tracing 19 | 20 | start_tracing() 21 | fn() 22 | stop_tracing() 23 | else: 24 | fn() 25 | end = time.perf_counter() 26 | 27 | print(json.dumps({"time": end - start})) 28 | -------------------------------------------------------------------------------- /lblprof/tests/example_scripts/complex_example.py: -------------------------------------------------------------------------------- 1 | import math 2 | import time 3 | import random 4 | from functools import lru_cache 5 | 6 | 7 | class Calculator: 8 | def __init__(self): 9 | self.values = [] 10 | 11 | def compute_square_roots(self, numbers): 12 | return [math.sqrt(n) for n in numbers] 13 | 14 | def simulate_workload(self): 15 | for _ in range(3): 16 | time.sleep(0.05) 17 | self.values.append(random.randint(1, 100)) 18 | 19 | 20 | @lru_cache(maxsize=None) 21 | def recursive_fib(n): 22 | if n <= 1: 23 | return n 24 | return recursive_fib(n - 1) + recursive_fib(n - 2) 25 | 26 | 27 | def generate_data(size): 28 | return [random.randint(1, 1000) for _ in range(size)] 29 | 30 | 31 | def main(): 32 | calc = Calculator() 33 | time.sleep(0.2) 34 | data = generate_data(5) 35 | calc.compute_square_roots(data) 36 | calc.simulate_workload() 37 | 38 | for i in range(10): 39 | recursive_fib(i) 40 | 41 | 42 | if __name__ == "__main__": 43 | main() 44 | -------------------------------------------------------------------------------- /lblprof/tests/example_scripts/chroma_vector_search.py: -------------------------------------------------------------------------------- 1 | # start_time = time.time() 2 | import chromadb 3 | import numpy as np 4 | from chromadb.config import Settings 5 | 6 | # 1. Set up ChromaDB in-memory client 7 | client = chromadb.Client(Settings(anonymized_telemetry=False)) 8 | 9 | # 2. Create collection 10 | collection = client.create_collection(name="demo_collection", get_or_create=True) 11 | 12 | # 3. Add a few random vectors (dimension = 5) 13 | np.random.seed(42) 14 | n_vectors = 5 15 | dim = 5 16 | 17 | for i in range(n_vectors): 18 | vector = np.random.rand(dim).tolist() 19 | collection.add( 20 | ids=[f"vec_{i}"], 21 | embeddings=[vector], 22 | metadatas=[{"index": i}], 23 | documents=[f"Random vector #{i}"], 24 | ) 25 | 26 | 27 | # 4. Create a random query vector 28 | query_vector = np.random.rand(dim).tolist() 29 | 30 | # 5. Perform similarity search 31 | results = collection.query( 32 | query_embeddings=[query_vector], 33 | n_results=3, 34 | include=["embeddings", "metadatas", "documents"], 35 | ) 36 | 37 | # print(f"Time taken: {time.time() - start_time} seconds") 38 | -------------------------------------------------------------------------------- /lblprof/__init__.py: -------------------------------------------------------------------------------- 1 | import logging 2 | import sys 3 | 4 | # Import the base tracer that works on all Python versions 5 | from .custom_sysmon import CodeMonitor 6 | 7 | 8 | __all__ = ["CodeMonitor"] 9 | # Create a singleton instance for the module 10 | tracer = CodeMonitor() 11 | 12 | # Import the sys.monitoring-based tracer if Python 3.12+ is available 13 | if hasattr(sys, "monitoring"): 14 | from .custom_sysmon import CodeMonitor 15 | 16 | __all__.append("CodeMonitor") 17 | tracer = CodeMonitor() 18 | else: 19 | logging.warning("Python 3.12+ is required to use the sys.monitoring-based tracer.") 20 | 21 | 22 | def start_tracing() -> None: 23 | """Start tracing code execution.""" 24 | tracer.start_tracing() 25 | 26 | 27 | def stop_tracing() -> None: 28 | """Stop tracing code execution.""" 29 | tracer.stop_tracing() 30 | tracer.tree.build_tree() 31 | 32 | 33 | def show_tree() -> None: 34 | """Display the tree structure.""" 35 | tracer.tree.display_tree() 36 | 37 | 38 | # Add a module-level function to expose the interactive UI 39 | def show_interactive_tree(min_time_s: float = 0.1): 40 | """Display an interactive tree in the terminal.""" 41 | tracer.tree.show_interactive(min_time_s=min_time_s) 42 | -------------------------------------------------------------------------------- /pyproject.toml: -------------------------------------------------------------------------------- 1 | [build-system] 2 | requires = ["setuptools>=61.0.0", "wheel"] 3 | build-backend = "setuptools.build_meta" 4 | 5 | [project] 6 | name = "lblprof" 7 | version = "0.1.8" 8 | description = "Line by line terminal based profiler" 9 | authors = [ 10 | { name="le-codeur-rapide", email="paul.vezia@gmail.com" } 11 | 12 | 13 | 14 | ] 15 | readme = "Readme.md" 16 | requires-python = ">=3.12" 17 | license = {text = "MIT"} 18 | classifiers = [ 19 | "Programming Language :: Python :: 3", 20 | "License :: OSI Approved :: MIT License", 21 | "Operating System :: OS Independent", 22 | ] 23 | dependencies = [ 24 | "pydantic>=2.0.0", 25 | ] 26 | 27 | [project.optional-dependencies] 28 | dev = [ 29 | "matplotlib>=3.10.1", 30 | "virtualenv>=20.30.0", 31 | "pandas>=2.0.3", 32 | "chromadb>=0.5.23", 33 | "ruff>=0.3.0", 34 | "pytest>=7.0.0", 35 | "black>=24.2.0", 36 | ] 37 | 38 | [tool.setuptools] 39 | packages = ["lblprof"] 40 | 41 | [tool.pytest.ini_options] 42 | testpaths = ["lblprof/tests"] 43 | python_files = ["test_*.py"] 44 | python_classes = ["Test*"] 45 | python_functions = ["test_*"] 46 | log_cli = true 47 | log_cli_level = "INFO" 48 | addopts = "-v -x" 49 | 50 | [tool.black] 51 | line-length = 88 52 | target-version = ["py312"] 53 | -------------------------------------------------------------------------------- /.github/workflows/mkdocs.yml: -------------------------------------------------------------------------------- 1 | name: Build GitHub Pages 2 | on: 3 | push: 4 | branches: 5 | - main 6 | paths: 7 | - 'docs/**' 8 | workflow_dispatch: 9 | permissions: 10 | contents: write 11 | pages: write 12 | id-token: write 13 | 14 | jobs: 15 | build_mkdocs: 16 | runs-on: ubuntu-latest 17 | 18 | steps: 19 | - name: Checkout code 20 | uses: actions/checkout@v4 21 | with: 22 | fetch-depth: 0 23 | - uses: actions/setup-python@v5 24 | with: 25 | python-version: 3.9 26 | - run: pip install \ 27 | mkdocs-material 28 | - run: mkdocs gh-deploy --config-file mkdocs.yml --force 29 | 30 | deploy_mkdocs: 31 | needs: build_mkdocs 32 | environment: 33 | name: github-pages 34 | url: ${{ steps.deployment.outputs.page_url }} 35 | runs-on: ubuntu-latest 36 | steps: 37 | - name: Checkout 38 | uses: actions/checkout@v4 39 | with: 40 | ref: gh-pages 41 | - name: Setup Pages 42 | uses: actions/configure-pages@v5 43 | - name: Upload artifact 44 | uses: actions/upload-pages-artifact@v3 45 | with: 46 | path: '.' 47 | - name: Deploy to GitHub Pages 48 | id: deployment 49 | uses: actions/deploy-pages@v4 -------------------------------------------------------------------------------- /lblprof/tests/demo_example_script.py: -------------------------------------------------------------------------------- 1 | import os 2 | import logging 3 | import sys 4 | import time 5 | import runpy 6 | 7 | sys.path.append(os.getcwd()) 8 | from lblprof import show_interactive_tree, start_tracing, stop_tracing, tracer 9 | 10 | logging.basicConfig(level=logging.DEBUG) 11 | 12 | 13 | path_example_folder = os.path.join(os.path.dirname(__file__), "example_scripts") 14 | script_name = "data_computation.py" 15 | # script_name = "chroma_vector_search.py" 16 | # script_name = "fibonacci.py" 17 | script_name = "import_pandas.py" 18 | # script_name = "list_comprehension.py" 19 | # script_name = "try_except.py" 20 | script_name = "double_zip.py" 21 | # script_name = "generator_function.py" 22 | script_path = os.path.join(path_example_folder, script_name) 23 | 24 | 25 | def main2(): 26 | time.sleep(1) 27 | 28 | 29 | def main(): 30 | time.sleep(1) 31 | main2() 32 | return 33 | 34 | 35 | # reset cache 36 | # start_time = time.perf_counter() 37 | # subprocess.run(["python", script_path]) 38 | # end_time = time.perf_counter() 39 | # print(f"Time taken: {end_time - start_time} seconds") 40 | 41 | 42 | # run the tracer for a bit and return the tree 43 | start_tracing() 44 | # time.sleep(1) 45 | # main() 46 | # import pandas as pd # noqa: E402,F401 47 | 48 | # Load and execute the example script 49 | start_time = time.perf_counter() 50 | runpy.run_path(script_path, run_name="__main__") 51 | end_time = time.perf_counter() 52 | print(f"Time taken: {end_time - start_time} seconds") 53 | 54 | stop_tracing() 55 | # print the tree 56 | # show_tree() 57 | show_interactive_tree(min_time_s=0.0) 58 | tracer.tree._save_events() 59 | tracer.tree._save_events_index() 60 | -------------------------------------------------------------------------------- /Readme.md: -------------------------------------------------------------------------------- 1 | # LBLProf 2 | 3 | LBLProf is simple a line by line terminal based time profiler. It allows you to track the duration of each line and get a tree of the execution of the code. 4 | 5 | It comes with a simple interactive terminal UI to navigate the tree and see the stats of each line. 6 | Example of the terminal ui: 7 | 8 | ![Example of the terminal ui](./docs/terminalui_showcase.png) 9 | 10 | > [!WARNING] 11 | > LBLProf is a tool that is based on the sys.monitoring API that is available in Python 3.12 and above. 12 | > **It means that you need to use Python 3.12 or above to use it.** 13 | 14 | # Documentation 15 | 16 | ## Installation 17 | 18 | ```bash 19 | pip install lblprof 20 | ``` 21 | 22 | The only dependency of this package is pydantic, the rest is standard library. 23 | 24 | 25 | ## Usage 26 | 27 | This package contains 4 main functions: 28 | - `start_tracing()`: Start the tracing of the code. 29 | - `stop_tracing()`: Stop the tracing of the code, build the tree and compute stats 30 | - `show_interactive_tree(min_time_s: float = 0.1)`: show the interactive duration tree in the terminal. 31 | - `show_tree()`: print the tree to console. 32 | 33 | ```python 34 | from lblprof import start_tracing, stop_tracing, show_interactive_tree, show_tree 35 | 36 | start_tracing() 37 | 38 | # Your code here (Any code) 39 | 40 | stop_tracing() 41 | show_tree() # print the tree to console 42 | show_interactive_tree() # show the interactive tree in the terminal 43 | ``` 44 | 45 | ## How it works 46 | 47 | LBLProf is based on the sys.monitoring API that is available in Python 3.12 and above. [PEP 669](https://peps.python.org/pep-0669/) 48 | 49 | This new API allow us to cut down tracing when we are entering installed package and limit the tracing to the code of the user. 50 | 51 | -------------------------------------------------------------------------------- /lblprof/benchmark/benchmark_overhead.py: -------------------------------------------------------------------------------- 1 | import subprocess 2 | import sys 3 | import json 4 | from typing import Literal 5 | from lblprof.benchmark.funcs import sample_functions 6 | import logging 7 | 8 | logging.basicConfig(level=logging.INFO, force=True) 9 | 10 | 11 | BenchRunMode = Literal["profiled", "unprofiled"] 12 | 13 | 14 | def run(module: str, fn: str, mode: BenchRunMode) -> float: 15 | """Run the isolated benchmark process and return the time taken.""" 16 | process_script_path = "lblprof/benchmark/bench_isolated_process.py" 17 | r = subprocess.check_output( 18 | [ 19 | sys.executable, 20 | process_script_path, 21 | module, 22 | fn, 23 | mode, 24 | ] 25 | ) 26 | return json.loads(r)["time"] 27 | 28 | 29 | def print_bench_result( 30 | unprofiled_times: list[float], profiled_times: list[float] 31 | ) -> None: 32 | unprofiled_time = sum(unprofiled_times) / len(unprofiled_times) 33 | profiled_time = sum(profiled_times) / len(profiled_times) 34 | overhead = profiled_time - unprofiled_time 35 | overhead_pct = ( 36 | (overhead / unprofiled_time) * 100 if unprofiled_time > 0 else float("inf") 37 | ) 38 | print(f"Unprofiled time: {unprofiled_time:.6f} s") 39 | print(f"Profiled time: {profiled_time:.6f} s") 40 | print(f"Overhead: {overhead:.6f} s ({overhead_pct:.2f}%)") 41 | print() 42 | 43 | 44 | for func in sample_functions: 45 | module = func.__module__ 46 | fn = func.__name__ 47 | number_of_runs = 10 48 | profiled_times: list[float] = [] 49 | unprofiled_times: list[float] = [] 50 | for _ in range(number_of_runs): 51 | t_profiled = run(module, fn, "profiled") 52 | t_unprofiled = run(module, fn, "unprofiled") 53 | profiled_times.append(t_profiled) 54 | unprofiled_times.append(t_unprofiled) 55 | 56 | print_bench_result(unprofiled_times=unprofiled_times, profiled_times=profiled_times) 57 | -------------------------------------------------------------------------------- /lblprof/tests/test_line_stats_tree.py: -------------------------------------------------------------------------------- 1 | import os 2 | import pytest 3 | from lblprof.line_stats_tree import LineStatsTree 4 | 5 | 6 | @pytest.fixture 7 | def tree(): 8 | """Fixture to create a LineStatsTree instance for each test.""" 9 | return LineStatsTree() 10 | 11 | 12 | def test_add_line_event(tree: LineStatsTree): 13 | """Test adding a line event to the tree.""" 14 | # Create a line with no parent (root line) 15 | tree.add_line_event( 16 | id=1, 17 | file_name="test_file.py", 18 | function_name="test_function", 19 | line_no=10, 20 | start_time=0.0, 21 | stack_trace=[("test_file.py", "test_function", 10)], 22 | ) 23 | 24 | assert len(tree.raw_events_list) == 1 25 | assert tree.raw_events_list[0]["id"] == 1 26 | assert tree.raw_events_list[0]["file_name"] == "test_file.py" 27 | assert tree.raw_events_list[0]["function_name"] == "test_function" 28 | assert tree.raw_events_list[0]["line_no"] == 10 29 | assert tree.raw_events_list[0]["start_time"] == 0.0 30 | assert tree.raw_events_list[0]["stack_trace"] == [ 31 | ("test_file.py", "test_function", 10) 32 | ] 33 | 34 | 35 | def test_add_line_event_with_parent(tree: LineStatsTree): 36 | """Test adding a line event with a parent in the tree.""" 37 | # Create a parent line 38 | tree.add_line_event( 39 | id=1, 40 | file_name="test_file.py", 41 | function_name="parent_function", 42 | line_no=5, 43 | start_time=0.0, 44 | stack_trace=[("test_file.py", "parent_function", 5)], 45 | ) 46 | 47 | # Create a child line 48 | tree.add_line_event( 49 | id=2, 50 | file_name="test_file.py", 51 | function_name="child_function", 52 | line_no=10, 53 | start_time=0.0, 54 | stack_trace=[("test_file.py", "child_function", 10)], 55 | ) 56 | 57 | # Check that both lines were added to the tree 58 | assert len(tree.raw_events_list) == 2 59 | 60 | 61 | def test_get_source_code(tree: LineStatsTree): 62 | """Test getting source code for a line.""" 63 | # Create a temporary file 64 | with open("temp_test_file.py", "w") as f: 65 | f.write("def test_function():\n print('Hello')\n") 66 | 67 | # Get source code 68 | source = tree._get_source_code("temp_test_file.py", 2) 69 | 70 | # Check that the source code is correct 71 | assert source == "print('Hello')" 72 | 73 | # Clean up 74 | os.remove("temp_test_file.py") 75 | -------------------------------------------------------------------------------- /lblprof/tests/test_tree_coherence.py: -------------------------------------------------------------------------------- 1 | import runpy 2 | import os 3 | import pytest 4 | import logging 5 | 6 | 7 | from lblprof.line_stats_tree import LineStatsTree 8 | from lblprof import show_tree, start_tracing, stop_tracing, tracer 9 | 10 | logging.basicConfig(level=logging.DEBUG) 11 | 12 | 13 | # It is hard to get reliable tests for some example of code, something that we 14 | # can do is to check that the tree is coherent 15 | # - A line should have a parent key if and only if it is a child of another line 16 | # - The sum of the time of the children should be equal to the child time of the parent 17 | # - Whatever the line, if it is time.sleep(n), then the time should be n 18 | 19 | # Get all Python files from the example_scripts directory 20 | EXAMPLE_SCRIPTS_DIR = os.path.join(os.path.dirname(__file__), "example_scripts") 21 | EXAMPLE_SCRIPTS = [ 22 | os.path.join(EXAMPLE_SCRIPTS_DIR, f) 23 | for f in os.listdir(EXAMPLE_SCRIPTS_DIR) 24 | if f.endswith(".py") and not f.startswith("__") 25 | ] 26 | 27 | 28 | @pytest.fixture(params=EXAMPLE_SCRIPTS, ids=lambda x: os.path.basename(x)) 29 | def tree(request): 30 | # run the tracer for a bit and return the tree 31 | start_tracing() 32 | 33 | # Load and execute the example script 34 | runpy.run_path(request.param, run_name="__main__") 35 | 36 | stop_tracing() 37 | # print the tree 38 | print(f"Tree for {os.path.basename(request.param)}:") 39 | show_tree() 40 | return tracer.tree 41 | 42 | 43 | def test_tree_coherence(tree: LineStatsTree): 44 | _validate_parent_child_relations(tree) 45 | _validate_time_sleep(tree) 46 | 47 | 48 | def _validate_parent_child_relations(tree: LineStatsTree): 49 | for line in tree.events_index.values(): 50 | if line.parent is None: 51 | continue 52 | assert ( 53 | line.parent in tree.events_index 54 | ), f"Parent key {line.parent} not found in tree: line {line}" 55 | assert line.id in [ 56 | child.id for child in tree.events_index[line.parent].childs.values() 57 | ], f"Line {line.id} should have parent key {line.parent}" 58 | 59 | 60 | def _validate_time_sleep(tree: LineStatsTree): 61 | for line in tree.root_lines: 62 | if "time.sleep" in line.source: 63 | n = line.source.split("time.sleep(")[1].split(")")[0] 64 | total_time = float(n) * line.hits * 1000 65 | assert line.time == pytest.approx( 66 | total_time, rel=0.1 67 | ), f"Line {line.id} should have time {total_time} but has time {line.time}" 68 | -------------------------------------------------------------------------------- /lblprof/line_stat_object.py: -------------------------------------------------------------------------------- 1 | from typing import List, Literal, NamedTuple, Tuple, Optional, Union 2 | 3 | from pydantic import BaseModel, Field, ConfigDict 4 | 5 | 6 | class LineKey(NamedTuple): 7 | file_name: str 8 | function_name: str 9 | line_no: Union[int, Literal["END_OF_FRAME"]] 10 | 11 | 12 | class LineEvent(NamedTuple): 13 | id: int 14 | file_name: str 15 | function_name: str 16 | line_no: Union[int, Literal["END_OF_FRAME"]] 17 | start_time: float 18 | stack_trace: list[Tuple[str, str, int]] 19 | 20 | 21 | class LineStats(BaseModel): 22 | """Statistics for a single line of code.""" 23 | 24 | model_config = ConfigDict(validate_assignment=True) 25 | 26 | id: int = Field(..., ge=0, description="Unique identifier for this line") 27 | 28 | # Key infos 29 | file_name: str = Field(..., min_length=1, description="File containing this line") 30 | function_name: str = Field( 31 | ..., min_length=1, description="Function containing this line" 32 | ) 33 | line_no: Union[int, Literal["END_OF_FRAME"]] = Field( 34 | ..., description="Line number in the source file" 35 | ) 36 | stack_trace: List[Tuple[str, str, int]] = Field( 37 | default_factory=list, description="Stack trace for this line" 38 | ) 39 | 40 | # Stats 41 | start_time: float = Field( 42 | ..., ge=0, description="Time when this line was first executed" 43 | ) 44 | hits: int = Field(..., ge=0, description="Number of times this line was executed") 45 | time: float = Field( 46 | description="Time spent on this line in milliseconds", default=0 47 | ) 48 | 49 | # Source code 50 | source: str = Field(..., min_length=1, description="Source code for this line") 51 | 52 | # Parent line that called this function 53 | # If None then it 54 | parent: Optional[int] = None 55 | 56 | # Children lines called by this line (populated during analysis) 57 | # We use a dict because it alows us to remove some childs in O(1) time 58 | # We need to remove children when we merge duplicated and when we remove END_OF_FRAME events 59 | childs: dict[int, "LineStats"] = Field(default_factory=dict) 60 | 61 | @property 62 | def event_id(self) -> int: 63 | """Get the unique id for this line.""" 64 | return self.id 65 | 66 | @property 67 | def event_key(self) -> Tuple[LineKey, Tuple[LineKey, ...]]: 68 | """Get the unique key for the event.""" 69 | return ( 70 | LineKey( 71 | file_name=self.file_name, 72 | function_name=self.function_name, 73 | line_no=self.line_no, 74 | ), 75 | tuple( 76 | LineKey(file_name=frame[0], function_name=frame[1], line_no=frame[2]) 77 | for frame in self.stack_trace 78 | ), 79 | ) 80 | -------------------------------------------------------------------------------- /lblprof/tests/example_scripts/data_computation.py: -------------------------------------------------------------------------------- 1 | import logging 2 | 3 | 4 | import random 5 | import requests 6 | import pandas as pd 7 | 8 | import numpy as np 9 | import time 10 | 11 | logging.basicConfig(level=logging.DEBUG) 12 | 13 | time.sleep(1) 14 | 15 | 16 | def fetch_exchange_rates(): 17 | logging.info("Fetching exchange rates from Frankfurter API...") 18 | url = "https://api.frankfurter.app/2024-12-01..2025-01-01" 19 | params = {"from": "USD", "to": "EUR,BRL"} 20 | r = requests.get(url, params=params) 21 | r.raise_for_status() 22 | data = r.json()["rates"] 23 | df = pd.DataFrame(data).T.sort_index() 24 | return df 25 | 26 | 27 | def fetch_bitcoin_prices(): 28 | logging.info("Fetching Bitcoin prices from CoinGecko...") 29 | url = "https://api.coingecko.com/api/v3/coins/bitcoin/market_chart" 30 | params = {"vs_currency": "usd", "days": "31", "interval": "daily"} # Approx 1 month 31 | r = requests.get(url, params=params) 32 | r.raise_for_status() 33 | prices = r.json()["prices"] 34 | df = pd.DataFrame(prices, columns=["timestamp", "btc_usd"]) 35 | df["date"] = pd.to_datetime(df["timestamp"], unit="ms").dt.date 36 | df = df.groupby("date").mean() 37 | df.index = pd.to_datetime(df.index) 38 | df = df[["btc_usd"]] 39 | return df 40 | 41 | 42 | def fetch_weather_data(): 43 | logging.info("Fetching weather data from Open-Meteo...") 44 | url = "https://api.open-meteo.com/v1/forecast" 45 | params = { 46 | "latitude": 48.8566, 47 | "longitude": 2.3522, 48 | "daily": "temperature_2m_max", 49 | "timezone": "Europe/Paris", 50 | "start_date": "2024-12-01", 51 | "end_date": "2025-01-01", 52 | } 53 | r = requests.get(url, params=params) 54 | r.raise_for_status() 55 | data = r.json() 56 | logging.info(f"data = {data}") 57 | df = pd.DataFrame( 58 | { 59 | "date": data["daily"]["time"], 60 | "temp_max": [ 61 | random.randint(10, 30) for _ in range(len(data["daily"]["time"])) 62 | ], 63 | } 64 | ) 65 | df["date"] = pd.to_datetime(df["date"]) 66 | df.set_index("date", inplace=True) 67 | return df 68 | 69 | 70 | def normalize_series(series): 71 | return (series - series.mean()) / series.std() 72 | 73 | 74 | def simulate_feature_extraction(df): 75 | df["btc_log"] = np.log(df["btc_usd"]) 76 | df["eur_change"] = df["EUR"].pct_change() 77 | logging.info(f"df = {df.head()}") 78 | df["temp_sin"] = np.sin(df["temp_max"] / 10) 79 | df["brl_rolling"] = df["BRL"].rolling(window=5).mean() 80 | df["interaction"] = df["btc_log"] * df["eur_change"] * df["temp_sin"] 81 | df.dropna(inplace=True) 82 | return df 83 | 84 | 85 | def main(): 86 | start = time.time() 87 | # exchange = fetch_exchange_rates() 88 | # print(f"timeaaa = {time.time() - start}") 89 | # btc = fetch_bitcoin_prices() 90 | # weather = fetch_weather_data() 91 | 92 | # df = exchange.join(btc).join(weather) 93 | 94 | # df = simulate_feature_extraction(df) 95 | 96 | # df[["btc_usd", "EUR", "BRL", "temp_max"]] = df[ 97 | # ["btc_usd", "EUR", "BRL", "temp_max"] 98 | # ].apply(normalize_series) 99 | 100 | # print(df.head()) 101 | 102 | # # Visualization 103 | # df[["btc_usd", "EUR", "BRL", "temp_max", "interaction"]].plot( 104 | # figsize=(12, 6), title="Normalized Time Series Data with Interaction" 105 | # ) 106 | 107 | end = time.time() 108 | print(f"\nTotal execution time: {end - start:.2f} seconds") 109 | return 110 | 111 | 112 | if __name__ == "__main__": 113 | main() 114 | # stop_tracing() 115 | # show_interactive_tree() 116 | -------------------------------------------------------------------------------- /lblprof/custom_sysmon.py: -------------------------------------------------------------------------------- 1 | import logging 2 | import sys 3 | import time 4 | from types import CodeType 5 | from typing import List, Tuple 6 | from .line_stats_tree import LineStatsTree 7 | 8 | 9 | # Check if sys.monitoring is available (Python 3.12+) 10 | if not hasattr(sys, "monitoring"): 11 | raise ImportError("sys.monitoring is not available. This requires Python 3.12+") 12 | 13 | 14 | class CodeMonitor: 15 | """ 16 | This class uses the sys.monitoring API (Python 3.12+) to register execution time for each line of code. 17 | It will create a monitore the infos of the execution time for each line of code, keeping track of the frame (the function / module) 18 | that called it. 19 | The goal of thsi class is to store the infos as fast as possible to have as few overhead as possible. All the computation and storage is 20 | done during the build_tree metho of the tree class 21 | """ 22 | 23 | def __init__(self): 24 | # Define a unique monitoring tool ID 25 | self.tool_id = sys.monitoring.PROFILER_ID 26 | 27 | # Call stack to store callers and keep track of the functions that called the current frame 28 | self.call_stack: List[Tuple[str, str, int]] = [] 29 | 30 | # Data structure to store the infos, insert should be quick during tracing, compute should be delayed at build time 31 | self.tree = LineStatsTree() 32 | 33 | # Use to store the line info until next line to keep track of who is the caller during call events 34 | self.tempo_line_infos: Tuple[str, str, int] | None = None 35 | 36 | # We use this to store the overhead of the monitoring tool and deduce it from the total time 37 | # The time of execution of the program will still be longer but at least the displayed time will be 38 | # more accurate 39 | self.overhead = 0 40 | 41 | # Total number of events, used to generate unique ids for each line 42 | self.total_events = 0 43 | 44 | def _handle_call(self, code: CodeType, instruction_offset: int): 45 | """Handle function call events""" 46 | start = time.perf_counter() 47 | file_name = code.co_filename 48 | 49 | if not self._is_user_code(file_name): 50 | # The call is from an imported module, we deactivate monitoring for this function 51 | if not sys.monitoring.get_tool(self.tool_id): 52 | sys.monitoring.use_tool_id(self.tool_id, "lblprof-monitor") 53 | sys.monitoring.set_local_events(self.tool_id, code, 0) 54 | self.overhead += time.perf_counter() - start 55 | return 56 | 57 | func_name = code.co_name 58 | line_no = code.co_firstlineno 59 | # We entered a frame from user code, we activate monitoring for it 60 | if not sys.monitoring.get_tool(self.tool_id): 61 | sys.monitoring.use_tool_id(self.tool_id, "lblprof-monitor") 62 | sys.monitoring.set_local_events( 63 | self.tool_id, 64 | code, 65 | sys.monitoring.events.LINE 66 | | sys.monitoring.events.PY_RETURN 67 | | sys.monitoring.events.PY_START, 68 | ) 69 | 70 | logging.debug( 71 | f"handle call: filename: {file_name}, func_name: {func_name}, line_no: {line_no}" 72 | ) 73 | if "" in func_name: 74 | return 75 | # We get info on who called the function 76 | # Using the tempo line infos instead of frame.f_back allows us to 77 | # get information about last parent that is from user code and not 78 | # from an imported / built in module 79 | if not self.tempo_line_infos: 80 | # Here we are called by a root line, so no caller in the stack 81 | return 82 | 83 | # caller_key is a tuple of (caller_file, caller_func, caller_line_no) 84 | caller_key = self.tempo_line_infos 85 | 86 | # Update call stack 87 | # Until we return from the function, all lines executed will have 88 | # the caller line as parent 89 | 90 | self.call_stack.append(caller_key) 91 | 92 | def _handle_line(self, code: CodeType, line_number: int): 93 | """Handle line execution events""" 94 | now = time.perf_counter() 95 | 96 | file_name = code.co_filename 97 | func_name = code.co_name 98 | line_no = line_number 99 | 100 | if not self._is_user_code(file_name): 101 | # The line is from an imported module, we deactivate monitoring for this line 102 | self.overhead += time.perf_counter() - now 103 | return 104 | 105 | # Add the line record to the tree 106 | logging.debug(f"tracing line: {file_name} {func_name} {line_no}") 107 | self.tree.add_line_event( 108 | id=self.total_events, 109 | file_name=file_name, 110 | function_name=func_name, 111 | line_no=line_no, 112 | # We substract the overhead to simulate a raw run without tracing 113 | start_time=now - self.overhead, 114 | stack_trace=self.call_stack.copy(), 115 | ) 116 | if line_no not in ["END_OF_FRAME", 0]: 117 | # In this case the line is not a real line of code so it can't be a parent to any other line 118 | self.tempo_line_infos = (file_name, func_name, line_no) 119 | self.total_events += 1 120 | 121 | def _handle_return(self, code: CodeType, instruction_offset: int, retval: object): 122 | """Handle function return events""" 123 | now = time.perf_counter() 124 | 125 | file_name = code.co_filename 126 | func_name = code.co_name 127 | line_no = code.co_firstlineno 128 | 129 | # Skip if not user code 130 | if not self._is_user_code(file_name): 131 | self.overhead += time.perf_counter() - now 132 | return sys.monitoring.DISABLE 133 | logging.debug(f"Returning from {func_name} in {file_name} ({line_no})") 134 | 135 | # In case the stop_tracing is called from a lower frame than start_tracing, 136 | # we need to activate monitoring for the returned frame 137 | current_frame = sys._getframe().f_back 138 | if not sys.monitoring.get_tool(self.tool_id): 139 | sys.monitoring.use_tool_id(self.tool_id, "lblprof-monitor") 140 | # We check if the returned frame already exists, if yes we activate monitoring for it 141 | if current_frame and current_frame.f_back and current_frame.f_back.f_code: 142 | sys.monitoring.set_local_events( 143 | self.tool_id, 144 | current_frame.f_code, 145 | sys.monitoring.events.LINE 146 | | sys.monitoring.events.PY_RETURN 147 | | sys.monitoring.events.PY_START, 148 | ) 149 | 150 | # Adding a END_OF_FRAME event to the tree to mark the end of the frame 151 | # This is used to compute the duration of the last line of the frame 152 | self.tree.add_line_event( 153 | id=self.total_events, 154 | file_name=file_name, 155 | function_name=func_name, 156 | line_no="END_OF_FRAME", 157 | start_time=now - self.overhead, 158 | stack_trace=self.call_stack.copy(), 159 | ) 160 | 161 | # A function is returning 162 | # We just need to pop the last line from the call stack so next 163 | # lines will have the correct parent 164 | if self.call_stack: 165 | self.call_stack.pop() 166 | 167 | self.total_events += 1 168 | 169 | def start_tracing(self) -> None: 170 | # Reset state 171 | self.__init__() 172 | if not sys.monitoring.get_tool(self.tool_id): 173 | sys.monitoring.use_tool_id(self.tool_id, "lblprof-monitor") 174 | else: 175 | # if the tool already assigned is not ours, we need to raise an error 176 | if sys.monitoring.get_tool(self.tool_id) != "lblprof-monitor": 177 | raise RuntimeError( 178 | "A tool with the id lblprof-monitor is already assigned, please stop it before starting a new tracing" 179 | ) 180 | 181 | # Register our callback functions 182 | sys.monitoring.register_callback( 183 | self.tool_id, sys.monitoring.events.PY_START, self._handle_call 184 | ) 185 | sys.monitoring.register_callback( 186 | self.tool_id, sys.monitoring.events.LINE, self._handle_line 187 | ) 188 | sys.monitoring.register_callback( 189 | self.tool_id, sys.monitoring.events.PY_RETURN, self._handle_return 190 | ) 191 | 192 | # 2 f_back to get out of the "start_tracing" function stack and get the user code frame 193 | current_frame = sys._getframe().f_back.f_back 194 | 195 | # The idea is that we register for calls at global level to not miss future calls and we register 196 | # for lines at the current frame (take care of set_local so it can be removed) 197 | sys.monitoring.set_events(self.tool_id, sys.monitoring.events.PY_START) 198 | sys.monitoring.set_local_events( 199 | self.tool_id, 200 | current_frame.f_code, 201 | sys.monitoring.events.LINE | sys.monitoring.events.PY_RETURN, 202 | ) 203 | logging.debug("Tracing started") 204 | 205 | def stop_tracing(self) -> None: 206 | # Turn off monitoring for our tool 207 | sys.monitoring.set_events(self.tool_id, 0) 208 | current_frame = sys._getframe() 209 | while current_frame: 210 | if not sys.monitoring.get_tool(self.tool_id): 211 | sys.monitoring.use_tool_id(self.tool_id, "lblprof-monitor") 212 | sys.monitoring.set_events(self.tool_id, 0) 213 | sys.monitoring.set_local_events(self.tool_id, current_frame.f_code, 0) 214 | sys.monitoring.free_tool_id(self.tool_id) 215 | current_frame = current_frame.f_back 216 | 217 | def _is_user_code(self, filename: str) -> bool: 218 | """Check if a file belongs to an installed module rather than user code. 219 | This is used to determine if we want to trace a line or not""" 220 | 221 | if ( 222 | ".local/lib" in filename 223 | or "/usr/lib" in filename 224 | or "/usr/local/lib" in filename 225 | or "site-packages" in filename 226 | or "dist-packages" in filename 227 | or "/lib/python3.12/" in filename 228 | or "frozen" in filename 229 | or ".local/share" in filename 230 | or "/.vscode-server/" in filename 231 | or filename.startswith("<") 232 | ): 233 | return False 234 | return True 235 | -------------------------------------------------------------------------------- /lblprof/curses_ui.py: -------------------------------------------------------------------------------- 1 | import curses 2 | from typing import Callable, List, Optional, TypedDict 3 | 4 | from lblprof.line_stat_object import LineStats 5 | 6 | 7 | class NodeTerminalUI(TypedDict): 8 | line: LineStats 9 | depth: int 10 | is_last: bool 11 | has_children: bool 12 | 13 | 14 | class TerminalTreeUI: 15 | """A terminal UI for displaying and interacting with tree data.""" 16 | 17 | def __init__( 18 | self, 19 | tree_data_provider: Callable[[Optional[LineStats]], List[LineStats]], 20 | node_formatter: Callable[[LineStats, str], str], 21 | ): 22 | """ 23 | Initialize the terminal UI. 24 | 25 | Args: 26 | tree_data_provider: Callable that returns the tree data to display 27 | node_formatter: Callable that formats a node for display 28 | """ 29 | self.tree_data_provider = tree_data_provider 30 | self.node_formatter = node_formatter 31 | 32 | # Tree branch characters 33 | self.branch_mid = "├── " 34 | self.branch_last = "└── " 35 | self.pipe = "│ " 36 | self.space = " " 37 | 38 | # UI state 39 | self.expanded_nodes = set() # Keys of expanded nodes (instead of collapsed) 40 | self.current_pos = 0 # Current selected position 41 | self.scroll_offset = 0 # Vertical scroll offset 42 | 43 | def _generate_display_data( 44 | self, root_nodes: List[LineStats] 45 | ) -> List[NodeTerminalUI]: 46 | """Generate flattened display data based on current UI state.""" 47 | display_data: List[NodeTerminalUI] = [] 48 | 49 | # Process each root node 50 | for i, root in enumerate(root_nodes): 51 | is_last_root = i == len(root_nodes) - 1 52 | 53 | node_data: NodeTerminalUI = { 54 | "line": root, 55 | "depth": 0, 56 | "is_last": is_last_root, 57 | "has_children": bool(root.childs), 58 | } 59 | display_data.append(node_data) 60 | 61 | # Add children only if explicitly expanded 62 | if root.event_key in self.expanded_nodes: 63 | self._add_children_to_display(display_data, root, 1) 64 | 65 | return display_data 66 | 67 | def _add_children_to_display( 68 | self, display_data: List[NodeTerminalUI], parent: LineStats, depth: int 69 | ): 70 | """Add children of a node to the display data recursively.""" 71 | # Get all child lines 72 | child_lines = self._get_sorted_children(parent) 73 | 74 | # Add each child to display data 75 | for i, child in enumerate(child_lines): 76 | is_last_child = i == len(child_lines) - 1 77 | 78 | node_data: NodeTerminalUI = { 79 | "line": child, 80 | "depth": depth, 81 | "is_last": is_last_child, 82 | "has_children": bool(child.childs), 83 | } 84 | display_data.append(node_data) 85 | 86 | # Add child's children only if explicitly expanded 87 | if child.event_key in self.expanded_nodes: 88 | self._add_children_to_display(display_data, child, depth + 1) 89 | 90 | def _get_sorted_children(self, parent: LineStats) -> List[LineStats]: 91 | """Get children of a parent node.""" 92 | # First get all valid children 93 | children = self.tree_data_provider(parent) 94 | 95 | # Group children by file 96 | children_by_file = {} 97 | for child in children: 98 | if child.file_name not in children_by_file: 99 | children_by_file[child.file_name] = [] 100 | children_by_file[child.file_name].append(child) 101 | 102 | # Sort each file's lines by line number 103 | for file_name in children_by_file: 104 | children_by_file[file_name].sort(key=lambda x: x.line_no) 105 | 106 | # Flatten all children 107 | all_children = [] 108 | for file_name in children_by_file: 109 | all_children.extend(children_by_file[file_name]) 110 | 111 | return all_children 112 | 113 | def _render_tree(self, stdscr, display_data: List[NodeTerminalUI], max_y, max_x): 114 | """Render the tree data on the screen.""" 115 | # Limit display data to visible area 116 | visible_height = max_y - 4 # Account for header and help 117 | 118 | # Initialize screen position and map of rows that need spacing 119 | screen_y = 2 # Start after header 120 | root_spacers = set() # Track where to add blank lines 121 | 122 | # Calculate visible range accounting for spacers 123 | visible_end = self.scroll_offset 124 | visible_items = 0 125 | for i in range(self.scroll_offset, len(display_data)): 126 | if visible_items >= visible_height: 127 | break 128 | visible_end = i + 1 129 | visible_items += 1 130 | if i + 1 in root_spacers: 131 | visible_items += 1 # Count the spacer 132 | 133 | # Render the visible portion 134 | rendered_pos = 0 135 | for i in range(self.scroll_offset, visible_end): 136 | # Add a blank line before root nodes (except the first one) 137 | if i in root_spacers: 138 | screen_y += 1 139 | rendered_pos += 1 140 | 141 | node = display_data[i] 142 | 143 | # Calculate screen position 144 | abs_pos = i 145 | 146 | # Color based on selection 147 | color = ( 148 | curses.color_pair(2) 149 | if abs_pos == self.current_pos 150 | else curses.color_pair(1) 151 | ) 152 | 153 | # Generate prefix based on depth and position 154 | prefix = self._get_prefix(node, display_data) 155 | 156 | # Get indicator for expandable nodes 157 | indicator = "" 158 | if node["has_children"]: 159 | indicator = ( 160 | "[+] " 161 | if node["line"].event_key not in self.expanded_nodes 162 | else "[-] " 163 | ) 164 | 165 | # Format the node using the provided formatter 166 | line_text = self.node_formatter(node["line"], indicator) 167 | 168 | # Combine and truncate if needed 169 | full_line = f"{prefix}{line_text}" 170 | if len(full_line) >= max_x: 171 | full_line = full_line[: max_x - 3] + "..." 172 | 173 | # Add to screen 174 | stdscr.addstr(screen_y, 0, full_line, color) 175 | screen_y += 1 176 | rendered_pos += 1 177 | 178 | def _get_prefix(self, node, display_data): 179 | """Generate the tree prefix for a node based on its position in the hierarchy.""" 180 | prefix = "" 181 | if node["depth"] == 0: 182 | # Root nodes 183 | return self.branch_last if node["is_last"] else self.branch_mid 184 | 185 | # For each level of depth, determine if we need a pipe or space 186 | for d in range(node["depth"] + 1): 187 | if d == node["depth"]: 188 | # Last level - add branch 189 | prefix += self.branch_last if node["is_last"] else self.branch_mid 190 | else: 191 | # Find if any parent at this level is a last child 192 | is_last_ancestor = self._check_if_last_ancestor(node, display_data, d) 193 | prefix += self.space if is_last_ancestor else self.pipe 194 | 195 | return prefix 196 | 197 | def _check_if_last_ancestor(self, node, display_data, depth): 198 | """Check if the node has an ancestor at the given depth that is the last child.""" 199 | # Find all nodes at this depth level 200 | nodes_at_depth = [n for n in display_data if n["depth"] == depth] 201 | if not nodes_at_depth: 202 | return False 203 | 204 | # Get the last node at this depth that appears before our target node 205 | for n in reversed(nodes_at_depth): 206 | if display_data.index(n) < display_data.index(node): 207 | return n["is_last"] 208 | 209 | return False 210 | 211 | def _toggle_collapse( 212 | self, display_data: List[NodeTerminalUI], current_node: NodeTerminalUI 213 | ): 214 | """Toggle collapse state of the current node.""" 215 | node_key = current_node["line"].event_key 216 | if node_key in self.expanded_nodes: 217 | self.expanded_nodes.remove(node_key) 218 | else: 219 | self.expanded_nodes.add(node_key) 220 | 221 | def run(self): 222 | """Run the terminal UI.""" 223 | curses.wrapper(self._main_curses_loop) 224 | 225 | def _main_curses_loop(self, stdscr): 226 | """Main curses loop for displaying and interacting with the tree.""" 227 | # Initialize curses 228 | stdscr.clear() 229 | curses.curs_set(0) # Hide cursor 230 | curses.start_color() 231 | curses.init_pair(1, curses.COLOR_WHITE, curses.COLOR_BLACK) # Normal text 232 | curses.init_pair(2, curses.COLOR_YELLOW, curses.COLOR_BLACK) # Highlighted text 233 | curses.init_pair(3, curses.COLOR_GREEN, curses.COLOR_BLACK) # Headers 234 | 235 | # Enable keypad and nodelay for better input handling 236 | stdscr.keypad(True) 237 | stdscr.nodelay(False) # Block and wait for input 238 | 239 | # Get root data from provider 240 | root_nodes = self.tree_data_provider(None) 241 | 242 | # Header and help text 243 | help_text = ( 244 | "[↑/↓]: Navigate | [PgUp/PgDn]: Page | [Enter]: Expand/Collapse | [q]: Quit" 245 | ) 246 | 247 | # Main UI loop 248 | running = True 249 | while running: 250 | # Get terminal dimensions 251 | max_y, max_x = stdscr.getmaxyx() 252 | 253 | # Adjust scroll_offset if window is resized to be smaller 254 | visible_height = max_y - 4 255 | display_data = self._generate_display_data(root_nodes) 256 | if self.current_pos >= len(display_data): 257 | self.current_pos = len(display_data) - 1 if display_data else 0 258 | 259 | # Make sure current position is visible 260 | if self.current_pos < self.scroll_offset: 261 | self.scroll_offset = self.current_pos 262 | elif self.current_pos >= self.scroll_offset + visible_height: 263 | self.scroll_offset = max(0, self.current_pos - visible_height + 1) 264 | 265 | # Clear screen 266 | stdscr.clear() 267 | 268 | # Display header 269 | header = "LINE TRACE TREE" 270 | stdscr.addstr(0, 0, header, curses.color_pair(3) | curses.A_BOLD) 271 | stdscr.addstr(1, 0, "=" * len(header), curses.color_pair(3)) 272 | 273 | # Display tree 274 | self._render_tree(stdscr, display_data, max_y, max_x) 275 | 276 | # Display help 277 | stdscr.addstr(max_y - 1, 0, help_text, curses.color_pair(3)) 278 | 279 | # Refresh screen 280 | stdscr.refresh() 281 | 282 | # Get user input 283 | key = stdscr.getch() 284 | 285 | # Handle key presses 286 | if key == curses.KEY_UP: 287 | # Move up 288 | if self.current_pos > 0: 289 | self.current_pos -= 1 290 | # Adjust scroll if needed 291 | if self.current_pos < self.scroll_offset: 292 | self.scroll_offset = self.current_pos 293 | 294 | elif key == curses.KEY_DOWN: 295 | # Move down 296 | if self.current_pos < len(display_data) - 1: 297 | self.current_pos += 1 298 | # Adjust scroll if needed 299 | visible_height = max_y - 4 300 | if self.current_pos >= self.scroll_offset + visible_height: 301 | self.scroll_offset = self.current_pos - visible_height + 1 302 | 303 | elif key == curses.KEY_PPAGE: # Page Up 304 | # Move up a page 305 | visible_height = max_y - 4 306 | self.current_pos = max(0, self.current_pos - visible_height) 307 | self.scroll_offset = max(0, self.scroll_offset - visible_height) 308 | 309 | elif key == curses.KEY_NPAGE: # Page Down 310 | # Move down a page 311 | visible_height = max_y - 4 312 | self.current_pos = min( 313 | len(display_data) - 1, self.current_pos + visible_height 314 | ) 315 | max_scroll = max(0, len(display_data) - visible_height) 316 | self.scroll_offset = min( 317 | max_scroll, self.scroll_offset + visible_height 318 | ) 319 | 320 | elif key == ord("\n"): # Enter key 321 | # Toggle collapse state 322 | if self.current_pos < len(display_data): 323 | current_node = display_data[self.current_pos] 324 | if current_node["has_children"]: 325 | self._toggle_collapse(display_data, current_node) 326 | 327 | elif key == ord("q"): # Quit 328 | running = False 329 | -------------------------------------------------------------------------------- /lblprof/line_stats_tree.py: -------------------------------------------------------------------------------- 1 | import logging 2 | import os 3 | from typing import List, Dict, Literal, Tuple, Optional, Union 4 | from lblprof.curses_ui import TerminalTreeUI 5 | from lblprof.line_stat_object import LineStats, LineKey, LineEvent 6 | 7 | 8 | class LineStatsTree: 9 | """A tree structure to manage LineStats objects with automatic parent-child time propagation.""" 10 | 11 | def __init__(self): 12 | # just raw event from custom_tracer 13 | # Inserting in this list should be as fast as possible to avoid overhead 14 | self.raw_events_list: List[LineEvent] = [] 15 | 16 | # Index of events by id 17 | self.events_index: Dict[int, LineStats] = {} 18 | 19 | # Track root nodes (lines of user code initial frame) 20 | self.root_lines: List[LineStats] = [] 21 | 22 | # cache of source code for lines 23 | # key is (file_name, line_no) for a source code line 24 | self.line_source: Dict[Tuple[str, int], str] = {} 25 | 26 | def add_line_event( 27 | self, 28 | id: int, 29 | file_name: str, 30 | function_name: str, 31 | line_no: Union[int, Literal["END_OF_FRAME"]], 32 | start_time: float, 33 | stack_trace: List[Tuple[str, str, int]], 34 | ) -> None: 35 | """Add a line event to the tree.""" 36 | # We don't want to add events from stop_tracing function 37 | # We might want something cleaner ? 38 | if "stop_tracing" in function_name: 39 | return 40 | 41 | logging.debug( 42 | f"Adding line event: {file_name}::{function_name}::{line_no} at {start_time}::{stack_trace}" 43 | ) 44 | 45 | self.raw_events_list.append( 46 | { 47 | "id": id, 48 | "file_name": file_name, 49 | "function_name": function_name, 50 | "line_no": line_no, 51 | "start_time": start_time, 52 | "stack_trace": stack_trace, 53 | } 54 | ) 55 | 56 | def build_tree(self) -> None: 57 | """Build the tree (self.events_index) from the raw events list.""" 58 | 59 | # 1. Build the events index (id: LineStats) 60 | for event in self.raw_events_list: 61 | source = self._get_source_code(event["file_name"], event["line_no"]) 62 | if "stop_tracing" in source: 63 | # This allow to delete the call to stop_tracing from the tree 64 | # and set the end line for the root lines 65 | event["line_no"] = "END_OF_FRAME" 66 | 67 | event_key = event["id"] 68 | if event_key not in self.events_index: 69 | self.events_index[event_key] = LineStats( 70 | id=event["id"], 71 | file_name=event["file_name"], 72 | function_name=event["function_name"], 73 | line_no=event["line_no"], 74 | stack_trace=event["stack_trace"], 75 | start_time=event["start_time"], 76 | hits=1, 77 | source=source, 78 | ) 79 | else: 80 | raise Exception("Event key already in self.events_index") 81 | 82 | # 2. Establish parent-child relationships 83 | # We first build a dict to map event keys to event ids, so we can get the parent_id in O(1) time for each line 84 | # we reverse the list because we always prefer that the parent of a line is the first event corresponding to the parent line 85 | linekey_to_id = { 86 | line.event_key[0]: event_id 87 | for event_id, line in reversed(list(self.events_index.items())) 88 | } 89 | for id, event in self.events_index.items(): 90 | 91 | # We get parent from stack trace 92 | if len(event.stack_trace) == 0: 93 | # we are in a root line 94 | self.root_lines.append(event) 95 | continue 96 | 97 | # find id of the parent in self.events_index 98 | parent_key = LineKey( 99 | file_name=event.stack_trace[-1][0], 100 | function_name=event.stack_trace[-1][1], 101 | line_no=event.stack_trace[-1][2], 102 | ) 103 | parent_id = linekey_to_id.get(parent_key) 104 | if parent_id is None: 105 | raise Exception( 106 | f"Parent key {event.stack_trace[-1]} not found in events index" 107 | ) 108 | 109 | self.events_index[parent_id].childs[id] = event 110 | event.parent = parent_id 111 | self.events_index[id] = event 112 | 113 | # 3. Update duration of each line 114 | # we use the time_save dict to store the id and start time of the previous line in the same frame (which is not necessary the previous line in the index) 115 | time_save: Dict[Union[int, None], Tuple[int, float]] = {} 116 | for id, event in self.events_index.items(): 117 | if event.parent not in time_save: 118 | # first line of the frame 119 | time_save[event.parent] = (id, event.start_time) 120 | continue 121 | # not the first line of the frame, update the time of the previous line 122 | previous_id, previous_start_time = time_save[event.parent] 123 | if event.start_time - previous_start_time < 0: 124 | logging.warning( 125 | f"Time of line {event.id} is negative: {event.start_time} - {previous_start_time}" 126 | ) 127 | self.events_index[previous_id].time = event.start_time - previous_start_time 128 | time_save[event.parent] = (id, event.start_time) 129 | 130 | # 4. Remove END_OF_FRAME lines 131 | # The END_OF_FRAME lines are not needed in the tree anymore 132 | for id, event in list(self.events_index.items()): 133 | if event.line_no == "END_OF_FRAME": 134 | parent_id = event.parent 135 | if parent_id is not None: 136 | del self.events_index[parent_id].childs[id] 137 | del self.events_index[id] 138 | self.root_lines = [ 139 | line for line in self.events_index.values() if line.parent is None 140 | ] 141 | 142 | # 5. Merge lines that have same file_name, function_name and line_no (to avoid duplicates in a for loop for example) 143 | # Not it is important to start by root nodes and merge going down the tree (DFS pre-order) 144 | grouped_events = {} 145 | 146 | def _merge(event: LineStats): 147 | """Merge events that have same file_name, function_name and line_no in the same frame.""" 148 | key = ( 149 | event.file_name, 150 | event.function_name, 151 | event.line_no, 152 | tuple(event.stack_trace), 153 | ) 154 | if key not in grouped_events: 155 | grouped_events[key] = event 156 | else: 157 | grouped = grouped_events[key] 158 | grouped.time += event.time 159 | grouped.hits += event.hits 160 | grouped.childs.update(event.childs) 161 | # update parent of the new children 162 | for child in event.childs.values(): 163 | child.parent = grouped.id 164 | 165 | # Now recurse on the children 166 | for child in event.childs.values(): 167 | _merge(child) 168 | 169 | for event in self.root_lines: 170 | _merge(event) 171 | 172 | # 6. Update the events_index with the merged events 173 | self.events_index = {} 174 | for key, event in grouped_events.items(): 175 | self.events_index[event.id] = event 176 | 177 | # 7. Update the childs attributes to remove deleted childs 178 | for id, event in self.events_index.items(): 179 | event.childs = { 180 | child.id: child 181 | for child in event.childs.values() 182 | if child.id in self.events_index 183 | } 184 | self.root_lines = [ 185 | line for line in self.events_index.values() if line.parent is None 186 | ] 187 | 188 | # -------------------------------- 189 | # Display methods 190 | # One method to print the tree in the console 191 | # One method to display the tree in an interactive terminal interface 192 | # -------------------------------- 193 | def display_tree( 194 | self, 195 | root_key: Optional[int] = None, 196 | depth: int = 0, 197 | max_depth: int = 10, 198 | is_last: bool = True, 199 | prefix: str = "", 200 | ) -> None: 201 | """Display a visual tree showing parent-child relationships between lines.""" 202 | if depth > max_depth: 203 | return # Prevent infinite recursion 204 | 205 | # Tree branch characters 206 | branch_mid = "├── " 207 | branch_last = "└── " 208 | pipe = "│ " 209 | space = " " 210 | 211 | def format_line_info(line: LineStats, branch: str): 212 | filename = os.path.basename(line.file_name) 213 | line_id = f"{filename}::{line.function_name}::{line.line_no}" 214 | 215 | # Truncate source code 216 | truncated_source = ( 217 | line.source[:60] + "..." if len(line.source) > 60 else line.source 218 | ) 219 | 220 | # Display line with time info and hits count 221 | assert line.time is not None 222 | return f"{prefix}{branch}{line_id} [hits:{line.hits} total:{line.time*1000:.2f}ms] - {truncated_source}" 223 | 224 | def group_children_by_file( 225 | children: dict[int, LineStats], 226 | ) -> Dict[str, List[LineStats]]: 227 | children_by_file = {} 228 | for child in children.values(): 229 | if child.file_name not in children_by_file: 230 | children_by_file[child.file_name] = [] 231 | children_by_file[child.file_name].append(child) 232 | 233 | # Sort each file's lines by line number 234 | for file_name in children_by_file: 235 | children_by_file[file_name].sort(key=lambda x: x.line_no) 236 | 237 | return children_by_file 238 | 239 | def get_all_children( 240 | children_by_file: Dict[str, List[LineStats]], 241 | ) -> List[LineStats]: 242 | all_children = [] 243 | for file_name in children_by_file: 244 | all_children.extend(children_by_file[file_name]) 245 | return all_children 246 | 247 | if root_key: 248 | 249 | line = self.events_index[root_key] 250 | branch = branch_last if is_last else branch_mid 251 | print(format_line_info(line, branch)) 252 | 253 | # Get all child lines 254 | child_lines = line.childs 255 | 256 | # Group and organize children 257 | children_by_file = group_children_by_file(child_lines) 258 | all_children = get_all_children(children_by_file) 259 | 260 | # Display child lines in order 261 | next_prefix = prefix + (space if is_last else pipe) 262 | for i, child in enumerate(all_children): 263 | is_last_child = i == len(all_children) - 1 264 | self.display_tree( 265 | child.id, depth + 1, max_depth, is_last_child, next_prefix 266 | ) 267 | else: 268 | # Print all root trees 269 | root_lines = self.root_lines 270 | if not root_lines: 271 | print("No root lines found in stats") 272 | return 273 | 274 | print("\n\nLINE TRACE TREE (HITS / SELF TIME / TOTAL TIME):") 275 | print("=================================================") 276 | 277 | # Sort roots by total time (descending) 278 | root_lines.sort( 279 | key=lambda x: x.line_no if x.line_no is not None else 0, reverse=False 280 | ) 281 | 282 | # For each root, render as a separate tree 283 | for i, root in enumerate(root_lines): 284 | is_last_root = i == len(root_lines) - 1 285 | branch = branch_last if is_last_root else branch_mid 286 | 287 | print(format_line_info(root, branch)) 288 | 289 | # Get all child lines and organize them 290 | children_by_file = group_children_by_file(root.childs) 291 | all_children = get_all_children(children_by_file) 292 | 293 | # Display child lines in order 294 | next_prefix = space if is_last_root else pipe 295 | for j, child in enumerate(all_children): 296 | is_last_child = j == len(all_children) - 1 297 | self.display_tree( 298 | child.id, 1, max_depth, is_last_child, next_prefix 299 | ) 300 | 301 | def show_interactive(self, min_time_s: float = 0.1): 302 | """Display the tree in an interactive terminal interface.""" 303 | 304 | # Define the data provider 305 | # Given a node (a line), return its children 306 | def get_tree_data(node_key: Optional[LineStats] = None) -> List[LineStats]: 307 | if node_key is None: 308 | # Return root nodes 309 | return [ 310 | line 311 | for line in self.root_lines 312 | if line.time and line.time >= min_time_s 313 | ] 314 | else: 315 | # Return children of the specified node 316 | return [ 317 | child 318 | for child in node_key.childs.values() 319 | if child.time and child.time >= min_time_s 320 | ] 321 | 322 | # Define the node formatter function 323 | # Given a line, return its formatted string (displayed in the UI) 324 | def format_node(line: LineStats, indicator: str = "") -> str: 325 | filename = os.path.basename(line.file_name) 326 | line_id = f"{filename}::{line.function_name}::{line.line_no}" 327 | 328 | # Truncate source code 329 | truncated_source = ( 330 | line.source[:40] + "..." if len(line.source) > 40 else line.source 331 | ) 332 | 333 | # Format stats 334 | assert line.time is not None 335 | stats = f"[hits:{line.hits} time:{line.time:.2f}s]" 336 | 337 | # Return formatted line 338 | return f"{indicator}{line_id} {stats} - {truncated_source}" 339 | 340 | # Create and run the UI 341 | ui = TerminalTreeUI(get_tree_data, format_node) 342 | ui.run() 343 | 344 | # -------------------------------- 345 | # Private methods 346 | # -------------------------------- 347 | def _save_events(self) -> None: 348 | """Save the events to a file.""" 349 | with open("events.csv", "w") as f: 350 | for event in self.raw_events_list: 351 | f.write( 352 | f"{event['id']},{event['file_name']},{event['function_name']},{event['line_no']},{event['start_time']},{event['stack_trace']}\n" 353 | ) 354 | 355 | def _save_events_index(self) -> None: 356 | """Save the events index to a file.""" 357 | with open("events_index.csv", "w") as f: 358 | for key, event in self.events_index.items(): 359 | f.write( 360 | f"{event.id},{event.file_name.split('/')[-1]},{event.function_name},{event.line_no},{event.source},{event.hits},{event.start_time},{event.time},{len(event.childs)},{event.parent}\n" 361 | ) 362 | 363 | def _get_source_code( 364 | self, file_name: str, line_no: Union[int, Literal["END_OF_FRAME"]] 365 | ) -> str: 366 | """Get the source code for a specific line in a file.""" 367 | if line_no == "END_OF_FRAME": 368 | return "END_OF_FRAME" 369 | if (file_name, line_no) in self.line_source: 370 | return self.line_source[(file_name, line_no)] 371 | try: 372 | with open(file_name, "r") as f: 373 | lines = f.readlines() 374 | if line_no - 1 < len(lines): 375 | source = lines[line_no - 1].strip() 376 | self.line_source[(file_name, line_no)] = source 377 | return source 378 | else: 379 | return " " 380 | except Exception as e: 381 | return ( 382 | f"Error getting source code of line {line_no} in file {file_name}: {e}" 383 | ) 384 | --------------------------------------------------------------------------------