refactor: type-hinting & PEP 561

Author: waldo121

.github/workflows/tests.yml

@@ -25,9 +25,12 @@ jobs:
       - name: Install dependencies
         run: |
           python -m pip install --upgrade pip
-          pip install pytest  
-          pip install .       
+          pip install . --group dev       
 
       - name: Run Pytest
         run: |
-          pytest tests/
+          pytest tests/
+
+      - name: Run Mypy
+        run: | 
+          mypy oracletrace

oracletrace/cli.py

@@ -2,15 +2,18 @@
 import sys
 import os
 import json
-import argparse
 import runpy
 import csv
 from .tracer import Tracer
 from .compare import compare_traces
+from typing import List, Dict, Any, Optional
+from re import Pattern
+from argparse import ArgumentParser, Namespace
+from pathlib import Path
 
 
-def main():
-    parser = argparse.ArgumentParser(
+def main() -> int:
+    parser: ArgumentParser  = ArgumentParser(
         description="OracleTrace - Lightweight execution tracer for Python projects"
     )
     parser.add_argument("target", help="Python script to trace")
@@ -39,21 +42,21 @@ def main():
         default=5.0,
         help="Regression threshold percentage used with --fail-on-regression.",
     )
-    args = parser.parse_args()
+    args: Namespace = parser.parse_args()
 
-    target = args.target
+    target: str = args.target
 
     if not os.path.exists(target):
         print(f"Target not found: {target}")
         return 1
 
     target = os.path.abspath(target)
-    root = os.getcwd()
-    target_dir = os.path.dirname(target)
+    root: str = os.getcwd()
+    target_dir: str = os.path.dirname(target)
     # Setup paths so imports work correctly in the target script
     sys.path.insert(0, target_dir)
-    ignored_args = [] if args.ignore is None else args.ignore
-    ignore_patterns = []
+    ignored_args: List[str] = [] if args.ignore is None else args.ignore
+    ignore_patterns: List[Pattern] = []
 
     for pattern in ignored_args:
         try:
@@ -63,14 +66,14 @@ def main():
             return 1
 
     # Start tracing, run the script, then stop
-    tracer = Tracer(root, ignore_patterns=ignore_patterns)
+    tracer: Tracer = Tracer(root, ignore_patterns=ignore_patterns)
     tracer.start()
     try:
         runpy.run_path(target, run_name="__main__")
     finally:
         tracer.stop()
 
-    data = tracer.get_trace_data()
+    data: Dict[str, Any] = tracer.get_trace_data()
 
     # Save json
     if args.json:
@@ -86,7 +89,7 @@ def main():
     # Export as csv
     if args.csv:
         with open(args.csv, "w", newline="", encoding="utf-8") as f:
-            writer = csv.DictWriter(f, fieldnames=["function", "total_time", "calls", "avg_time"])
+            writer: csv.DictWriter = csv.DictWriter(f, fieldnames=["function", "total_time", "calls", "avg_time"])
             writer.writeheader()
             for fn in data["functions"]:
                 writer.writerow({
@@ -96,7 +99,7 @@ def main():
                     "avg_time":   fn["avg_time"],
                 })
 
-    comparison_result = None
+    comparison_result: Optional[Dict[str,Any]] = None
 
     # Compare jsons
     if args.compare:
@@ -105,7 +108,7 @@ def main():
             return 1
 
         with open(args.compare, "r", encoding="utf-8") as f:
-            old_data = json.load(f)
+            old_data: Dict[str, Any] = json.load(f)
 
         comparison_result = compare_traces(old_data, data, threshold=args.threshold)
 

oracletrace/compare.py

@@ -1,19 +1,20 @@
 from rich import print
+from typing import Any, Dict, List, Set, Optional
 
 
-def compare_traces(old_data, new_data, threshold=5.0):
-    old_funcs = {f["name"]: f for f in old_data["functions"]}
-    new_funcs = {f["name"]: f for f in new_data["functions"]}
+def compare_traces(old_data: Dict[str, Any], new_data: Dict[str, Any], threshold: float = 5.0) -> Dict[str, Any]:
+    old_funcs: Dict[str, Any] = {f["name"]: f for f in old_data["functions"]}
+    new_funcs: Dict[str, Any] = {f["name"]: f for f in new_data["functions"]}
 
-    regressions = []
+    regressions: List[Dict[str, Any]] = []
 
     print("\n[bold cyan]Comparison Results:[/]\n")
 
-    all_functions = set(old_funcs) | set(new_funcs)
+    all_functions: Set[str] = set(old_funcs) | set(new_funcs)
 
     for name in sorted(all_functions):
-        old = old_funcs.get(name)
-        new = new_funcs.get(name)
+        old: Optional[None] = old_funcs.get(name)
+        new: Optional[None] = new_funcs.get(name)
 
         if not old:
             print(f"[green]+ {name} (new function)[/]")
@@ -23,16 +24,16 @@ def compare_traces(old_data, new_data, threshold=5.0):
             print(f"[red]- {name} (removed)[/]")
             continue
 
-        old_time = old["total_time"]
-        new_time = new["total_time"]
+        old_time: float = old["total_time"]
+        new_time: float = new["total_time"]
 
         if old_time == 0:
             continue
 
-        diff = new_time - old_time
-        percent = (diff / old_time) * 100
+        diff: float = new_time - old_time
+        percent: float = (diff / old_time) * 100
 
-        color = "red" if percent > threshold else "green" if percent < -threshold else "yellow"
+        color: str = "red" if percent > threshold else "green" if percent < -threshold else "yellow"
 
         print(
             f"{name}\n"

oracletrace/py.typed

@@ -5,36 +5,40 @@
 from rich.tree import Tree
 from rich.table import Table
 from rich import print
+from typing import List, Optional, Callable, DefaultDict, Any, Tuple, Dict
+from re import Pattern
+from pathlib import Path
+from types import FrameType
 
 
 class Tracer:
-    def __init__(self, root_dir, ignore_patterns = None):
-        self._root_path = os.path.abspath(root_dir)
-        self._call_stack = []
-        self._func_calls = defaultdict(int)
-        self._func_time = defaultdict(float)
-        self._call_map = defaultdict(lambda: defaultdict(int))
-        self._original_profile_func = None
-        self._enabled = False
-        self._start_time = 0.0          
-        self._total_time = 0.0
-        self._ignore_patterns = ignore_patterns
-
-
-    def start(self):
+    def __init__(self, root_dir: str, ignore_patterns: Optional[List[Pattern]] = None) -> None:
+        self._root_path: str = os.path.abspath(root_dir)
+        self._call_stack: List[Tuple[int, str, float]] = []
+        self._func_calls: DefaultDict[str, int] = defaultdict(int)
+        self._func_time: DefaultDict[str, float] = defaultdict(float)
+        self._call_map: DefaultDict[str, DefaultDict[str, int]] = defaultdict(lambda: defaultdict(int))
+        self._original_profile_func: Optional[Callable[[FrameType, str, Any], object]] = None
+        self._enabled: bool = False
+        self._start_time: float = 0.0          
+        self._total_time: float = 0.0
+        self._ignore_patterns: Optional[List[Pattern]] = ignore_patterns
+
+
+    def start(self) -> None:
         # Start Tracer
         self._enabled = True
         self._start_time = time.perf_counter()          
         self._original_profile_func = sys.getprofile()
         sys.setprofile(self._trace)
 
-    def stop(self):
+    def stop(self) -> None:
         # Stops Tracer
         self._enabled = False
         self._total_time = time.perf_counter() - self._start_time  
         sys.setprofile(self._original_profile_func)
 
-    def _is_ignored(self, filename):
+    def _is_ignored(self, filename: str) -> bool:
         # Return true if the filename should be ignored
         if not self._ignore_patterns:
             return False
@@ -45,46 +49,46 @@ def _is_ignored(self, filename):
 
         return False
 
-    def _is_user_code(self, filename):
+    def _is_user_code(self, filename: str) -> bool:
         # Filter out files not in the project root
-        if not filename.startswith(self._root_path):
+        if not filename.startswith(str(self._root_path)):
             return False
         # Filter out third-party libraries
         if "site-packages" in filename or "dist-packages" in filename:
             return False
         return True
 
-    def _get_key(self, frame):
-        co_filename = frame.f_code.co_filename
+    def _get_key(self, frame: FrameType) -> Optional[str]:
+        co_filename: str = frame.f_code.co_filename
         # Ignore internal python frames (e.g. <string>)
         if co_filename.startswith("<"):
             return None
-        filename = os.path.abspath(co_filename)
+        filename: str = os.path.abspath(co_filename)
         # Check if the file belongs to the user's project
         if not self._is_user_code(filename):
             return None
         # Create a relative path key for readability
-        rel_path = os.path.relpath(filename, self._root_path)
-        qualname = getattr(frame.f_code, "co_qualname", frame.f_code.co_name)
-        key = f"{rel_path}:{qualname}"
+        rel_path: str = os.path.relpath(filename, self._root_path)
+        qualname: str = getattr(frame.f_code, "co_qualname", frame.f_code.co_name)
+        key: str = f"{rel_path}:{qualname}"
 
         # Check if the file should be ignored based on inputted ignoring pattern
         if self._is_ignored(key):
             return None
 
         return key
 
-    def _trace(self, frame, event, arg):
+    def _trace(self, frame: FrameType, event: str, _: Any) -> None:
         try:
             if not self._enabled:
                 return
 
             if event == "call":
-                key = self._get_key(frame)
+                key: Optional[str] = self._get_key(frame)
                 if not key:
                     return
 
-                caller = self._call_stack[-1][1] if self._call_stack else "<module>"
+                caller: str = self._call_stack[-1][1] if self._call_stack else "<module>"
                 self._call_map[caller][key] += 1
                 self._func_calls[key] += 1
                 self._call_stack.append((id(frame), key, time.perf_counter()))
@@ -99,8 +103,8 @@ def _trace(self, frame, event, arg):
                     self._func_time[key] += time.perf_counter() - start
                 else:
                     # Stack unwinding (handle exceptions or missed returns)
-                    fid = id(frame)
-                    found = False
+                    fid: int = id(frame)
+                    found: bool = False
                     # Search for the frame in the stack from top to bottom
                     for i in range(len(self._call_stack) - 1, -1, -1):
                         if self._call_stack[i][0] == fid:
@@ -117,15 +121,15 @@ def _trace(self, frame, event, arg):
         except Exception as e:
             print(f"[bold red]Error in oracletrace tracer: {e}[/]", file=sys.stderr)
 
-    def show_results(self, _top):
+    def show_results(self, _top: Optional[int]) -> None:
         if not self._func_calls:
             print("[yellow]No calls traced.[/]")
             return
 
         # Summary table
         print("[bold green]Summary:[/]")
         print(f"[bold cyan]Total execution time: {self._total_time:.4f}s[/]")
-        table = Table(title="Top functions by Total Time")
+        table: Table = Table(title="Top functions by Total Time")
         table.add_column("Function", justify="left", style="cyan", no_wrap=True)
         table.add_column("Total Time (s)", justify="right", style="magenta")
         table.add_column("Calls", justify="right", style="green")
@@ -145,11 +149,11 @@ def show_results(self, _top):
 
         print("\n[bold green]Logic Flow:[/]")
 
-        tree = Tree("[bold yellow]<module>[/]")
+        tree: Tree = Tree("[bold yellow]<module>[/]")
 
         # Recursively build the execution tree
-        def add_nodes(parent_node, parent_key, current_path):
-            children = self._call_map.get(parent_key, {})
+        def add_nodes(parent_node: Tree, parent_key: str, current_path: set[str]) -> None:
+            children: DefaultDict[str,int] = self._call_map.get(parent_key, defaultdict(int))
             # Sort children by total execution time
             sorted_children = sorted(
                 children.items(),
@@ -171,8 +175,8 @@ def add_nodes(parent_node, parent_key, current_path):
         add_nodes(tree, "<module>", {"<module>"})
         print(tree)
 
-    def get_trace_data(self):
-        functions = []
+    def get_trace_data(self) -> Dict[str, Any]:
+        functions: List[Dict[str, Any]] = []
 
         for key, total_time in self._func_time.items():
             calls = self._func_calls[key]
@@ -198,10 +202,10 @@ def get_trace_data(self):
         }
 
 
-_tracer_instance = None
+_tracer_instance: Optional[Tracer] = None
 
 
-def start_trace(root_dir):
+def start_trace(root_dir: str) -> None:
     # Starts tracer instance
     global _tracer_instance
     if _tracer_instance is not None:
@@ -211,7 +215,7 @@ def start_trace(root_dir):
     _tracer_instance.start()
 
 
-def stop_trace():
+def stop_trace() -> Optional[Dict[str,Any]]:
     # Stops tracer instance
     global _tracer_instance
     if _tracer_instance:
@@ -224,10 +228,10 @@ def stop_trace():
         return None
 
 
-def show_results():
+def show_results(top: Optional[int]) -> None:
     # Show results from global tracer instance
     global _tracer_instance
     if _tracer_instance:
-        _tracer_instance.show_results()
+        _tracer_instance.show_results(top)
     else:
         print("[yellow]Tracer was not started.[/]")