11import json
22import os
3+ from typing import Any , TypedDict
34
45import click
56import numpy as np
7+ from numpy .typing import NDArray
68
79BIG_WIDTH = 80
810SMALL_WIDTH = 8
911
1012
11- def find_json_files (base_dir ):
13+ class TEST_STATS_TYPE (TypedDict ):
14+ durations : list [str | float ]
15+ n_tests : int
16+
17+
18+ def find_json_files (base_dir : str ) -> list [str ]:
1219 """Recursively find all JSON files in subdirectories."""
13- json_files = []
20+ json_files : list [ str ] = []
1421 for root , _ , files in os .walk (base_dir ):
1522 for file in files :
1623 if file .endswith (".jsonl" ):
1724 json_files .append (os .path .join (root , file ))
1825 return json_files
1926
2027
21- def read_json_file (file_path ) :
28+ def read_json_file (file_path : str ) -> list [ dict [ str , str ]] :
2229 """Read a JSON file and return its content as a list of test configurations."""
2330 with open (file_path , "r" , encoding = "utf-8" ) as f :
2431 try :
@@ -29,24 +36,28 @@ def read_json_file(file_path):
2936 return []
3037
3138
32- def extract_tests_with_tags (json_files ) :
39+ def extract_tests_with_tags (json_files : list [ str ]) -> list [ dict [ str , str | list [ str ]]] :
3340 """Extract test data and assign a tag based on the directory name."""
34- tests = []
41+ tests : list [ dict [ str , str | list [ str ]]] = []
3542
3643 for file_path in json_files :
3744 directory_name = os .path .basename (os .path .dirname (file_path ))
3845 test_data = read_json_file (file_path )
3946
4047 for test in test_data :
4148 if test .get ("outcome" , "" ).lower () == "passed" and test .get ("duration" ):
42- nodeid = test .get ("nodeid" )
49+ nodeid : str = test .get ("nodeid" , "" )
50+
4351 if nodeid .startswith ("tests/" ):
4452 nodeid = nodeid [6 :]
4553
46- when = test .get ("when" )
47- duration = test ["duration" ]
48- tags = directory_name .split ("-" )
49- tags .remove ("logs" )
54+ when : str = test .get ("when" , "" )
55+ duration : str = test ["duration" ]
56+ tags : list [str ] = directory_name .split ("-" )
57+
58+ if "logs" in tags :
59+ tags .remove ("logs" )
60+
5061 id_ = f"{ nodeid } { when }
5162
5263 tests .append (
@@ -61,12 +72,14 @@ def extract_tests_with_tags(json_files):
6172 return tests
6273
6374
64- def compute_statistics (tests ):
75+ def compute_statistics (
76+ tests : list [dict [str , str | list [str ]]],
77+ ) -> list [dict [str , str | float ]]:
6578 """Compute average duration and standard deviation per test ID."""
66- test_stats = {}
79+ test_stats : dict [ str , TEST_STATS_TYPE ] = {}
6780
6881 for test in tests :
69- test_id = test ["id" ]
82+ test_id : str = test ["id" ]
7083 if test_id not in test_stats :
7184 test_stats [test_id ] = {
7285 "durations" : [],
@@ -76,10 +89,10 @@ def compute_statistics(tests):
7689 test_stats [test_id ]["durations" ].append (test ["duration" ])
7790 test_stats [test_id ]["n_tests" ] += 1
7891
79- summary = []
92+ summary : list [ dict [ str , Any ]] = []
8093
8194 for test_id , data in test_stats .items ():
82- durations = np .array (data ["durations" ])
95+ durations : NDArray [ Any ] = np .array (data ["durations" ])
8396
8497 if durations .size == 0 :
8598 continue
@@ -119,10 +132,15 @@ def compute_statistics(tests):
119132 return summary
120133
121134
122- def print_table (data , keys , headers , title = "" ):
135+ def print_table (
136+ data : list [dict [str , str | float ]],
137+ keys : list [str ],
138+ headers : list [str ],
139+ title : str = "" ,
140+ ):
123141 JUNCTION = "|"
124142
125- def make_bold (s ) :
143+ def make_bold (s : str ) -> str :
126144 return click .style (s , bold = True )
127145
128146 h = [headers [0 ].ljust (BIG_WIDTH )]
@@ -135,7 +153,7 @@ def make_bold(s):
135153 + f"-{ JUNCTION } .join (["-" * len (each ) for each in h ])
136154 + f"-{ JUNCTION }
137155 )
138- top_sep = f"{ JUNCTION } + "-" * (len_h - 2 ) + f"{ JUNCTION }
156+ # top_sep: str = f"{JUNCTION}" + "-" * (len_h - 2) + f"{JUNCTION}"
139157
140158 if title :
141159 # click.echo(top_sep)
@@ -148,17 +166,17 @@ def make_bold(s):
148166 click .echo (sep )
149167
150168 for test in data :
151- s = []
169+ s : list [ str ] = []
152170 for i , each_key in enumerate (keys ):
153171
154172 if i == 0 :
155- id_ = test [each_key ]
173+ id_ : str = test [each_key ]
156174
157175 id_ = (
158- id_ .replace ("(" , "\ ("
159- .replace (")" , "\ )"
160- .replace ("[" , "\ ["
161- .replace ("]" , "\ ]"
176+ id_ .replace ("(" , r" ("
177+ .replace (")" , r" )"
178+ .replace ("[" , r" ["
179+ .replace ("]" , r" ]"
162180 )
163181 if len (id_ ) >= BIG_WIDTH :
164182 id_ = id_ [: BIG_WIDTH - 15 ] + "..." + id_ [- 12 :]
@@ -177,7 +195,7 @@ def make_bold(s):
177195 # click.echo(sep)
178196
179197
180- def print_summary (summary , num = 10 ):
198+ def print_summary (summary : list [ dict [ str , str | float ]], num : int = 10 ):
181199 """Print the top N longest tests and the top N most variable tests."""
182200 longest_tests = sorted (summary , key = lambda x : - x ["average_duration" ])[:num ]
183201 most_variable_tests = sorted (summary , key = lambda x : - x ["std_dev" ])[:num ]
@@ -225,15 +243,20 @@ def print_summary(summary, num=10):
225243 default = None ,
226244)
227245@click .option (
228- "--num" , default = 10 , help = "Number of top tests to display." , show_default = True
246+ "--num" ,
247+ type = int ,
248+ default = 10 ,
249+ help = "Number of top tests to display." ,
250+ show_default = True ,
229251)
230252@click .option (
231253 "--save-file" ,
232254 default = None ,
255+ type = click .Path (exists = False , dir_okay = False ),
233256 help = "File to save the test durations. Default 'tests_durations.json'." ,
234257 show_default = True ,
235258)
236- def analyze_tests (directory , num , save_file ):
259+ def analyze_tests (directory : str , num : int , save_file : str ):
237260 directory = directory or os .getcwd () # Change this to your base directory
238261 json_files = find_json_files (directory )
239262 tests = extract_tests_with_tags (json_files )
0 commit comments