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,19 @@
 import sys
 import os
 import json
-import argparse
 import runpy
 import csv
-from .tracer import Tracer
-from .compare import compare_traces
+from .tracer import Tracer, TracerData
+from .compare import compare_traces, ComparisonData
+from typing import List, Dict, Any, Optional
+from re import Pattern
+from argparse import ArgumentParser, Namespace
+from pathlib import Path
+from dataclasses import asdict
 
 
-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 +43,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,19 +67,19 @@ 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: TracerData = tracer.get_trace_data()
 
     # Save json
     if args.json:
         with open(args.json, "w", encoding="utf-8") as f:
-            json.dump(data, f, indent=4)
+            json.dump(asdict(data), f, indent=4)
 
     # Display the analysis
     if args.top:
@@ -86,17 +90,17 @@ 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"]:
+            for fn in data.functions:
                 writer.writerow({
-                    "function":   fn["name"],
-                    "total_time": fn["total_time"],
-                    "calls":      fn["call_count"],
-                    "avg_time":   fn["avg_time"],
+                    "function":   fn.name,
+                    "total_time": fn.total_time,
+                    "calls":      fn.call_count,
+                    "avg_time":   fn.avg_time,
                 })
 
-    comparison_result = None
+    comparison_result: Optional[ComparisonData] = None
 
     # Compare jsons
     if args.compare:
@@ -105,11 +109,11 @@ def main():
             return 1
 
         with open(args.compare, "r", encoding="utf-8") as f:
-            old_data = json.load(f)
+            old_data: TracerData = TracerData.from_dict(json.load(f))
 
         comparison_result = compare_traces(old_data, data, threshold=args.threshold)
 
-        if args.fail_on_regression and comparison_result["has_regression"]:
+        if args.fail_on_regression and comparison_result.has_regression:
             print(
             f"Build failed: performance regression above {args.threshold:.2f}% detected."
             )

oracletrace/compare.py

@@ -1,19 +1,33 @@
+from .tracer import TracerData, FunctionData
 from rich import print
+from typing import Any, Dict, List, Set, Optional
+from dataclasses import dataclass
 
+@dataclass
+class RegressionData:
+    name: str
+    old_time: float
+    new_time: float
+    percent: float
 
-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"]}
+@dataclass
+class ComparisonData:
+    regressions: List[RegressionData]
+    has_regression: bool
 
-    regressions = []
+def compare_traces(old_data: TracerData, new_data: TracerData, threshold: float = 5.0) -> ComparisonData:
+    old_funcs: Dict[str, FunctionData] = {f.name: f for f in old_data.functions}
+    new_funcs: Dict[str, FunctionData] = {f.name: f for f in new_data.functions}
+
+    regressions: List[RegressionData] = []
 
     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[FunctionData] = old_funcs.get(name)
+        new: Optional[FunctionData] = new_funcs.get(name)
 
         if not old:
             print(f"[green]+ {name} (new function)[/]")
@@ -23,16 +37,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"
@@ -42,15 +56,15 @@ def compare_traces(old_data, new_data, threshold=5.0):
 
         if percent > threshold:
             regressions.append(
-                {
-                    "name": name,
-                    "old_time": old_time,
-                    "new_time": new_time,
-                    "percent": percent,
-                }
+                RegressionData(
+                        name = name,
+                        new_time = new_time,
+                        old_time = old_time,
+                        percent = percent
+                    )
             )
 
-    return {
-        "regressions": regressions,
-        "has_regression": len(regressions) > 0,
-    }
+    return ComparisonData(
+        regressions = regressions,
+        has_regression = len(regressions) > 0
+    )