99import tempfile
1010import subprocess
1111import sys
12+ import re
13+ import random
1214
1315from collections import defaultdict
1416
@@ -25,6 +27,45 @@ class WheelInfo:
2527 abi_tag : str
2628 platform_tag : str
2729
30+ def _python_tag_sort_key (tag : str ):
31+ key_parts = []
32+ for component in tag .split ("." ):
33+ match = re .match (r"([a-zA-Z]+)(\d+)$" , component )
34+ if match :
35+ prefix , numeric = match .groups ()
36+ key_parts .append ((prefix , int (numeric )))
37+ else :
38+ key_parts .append ((component , 0 ))
39+ return tuple (key_parts )
40+
41+ def _run_self_test ():
42+ logger .info ("Running combine_wheels self-test" )
43+ simple_tags = ["cp39" , "cp310" , "cp311" , "cp38" ]
44+ expected_simple = ["cp38" , "cp39" , "cp310" , "cp311" ]
45+ actual_simple = sorted (simple_tags , key = _python_tag_sort_key )
46+ assert actual_simple == expected_simple , f"Simple ordering failed: { actual_simple } "
47+
48+ composite_tags = ["cp39.cp310" , "cp39.cp312" , "cp38.cp310" , "cp310.cp311" , "cp39.cp311" ]
49+ expected_composite = ["cp38.cp310" , "cp39.cp310" , "cp39.cp311" , "cp39.cp312" , "cp310.cp311" ]
50+ actual_composite = sorted (composite_tags , key = _python_tag_sort_key )
51+ assert actual_composite == expected_composite , f"Composite ordering failed: { actual_composite } "
52+
53+ wheels = [
54+ WheelInfo (
55+ wheel_name = f"depthai-foo-{ tag } -abi-{ idx } " ,
56+ wheel_path = f"/tmp/{ tag } -{ idx } .whl" ,
57+ wheel_dvb = "depthai-foo" ,
58+ python_tag = tag ,
59+ abi_tag = f"abi{ idx } " ,
60+ platform_tag = "plat" ,
61+ )
62+ for idx , tag in enumerate (composite_tags )
63+ ]
64+ random .shuffle (wheels )
65+ wheels .sort (key = lambda info : _python_tag_sort_key (info .python_tag ))
66+ assert [w .python_tag for w in wheels ] == expected_composite , "WheelInfo sorting failed"
67+ logger .info ("Self-test passed" )
68+
2869def combine_wheels_linux (args , wheel_infos ):
2970
3071 logger .info ("Combining wheels for Linux!" )
@@ -295,6 +336,11 @@ def main(args: argparse.Namespace):
295336 format = FORMAT ,
296337 datefmt = DATEFMT
297338 )
339+ if args .self_test :
340+ _run_self_test ()
341+ return
342+ if not args .input_folder :
343+ raise ValueError ("--input_folder is required unless --self-test is specified" )
298344 ## Get a list of all wheels in the input folder
299345 wheels = glob .glob (os .path .join (args .input_folder , "*.whl" ))
300346
@@ -315,7 +361,8 @@ def main(args: argparse.Namespace):
315361 abi_tag = abi_tag ,
316362 platform_tag = platform_tag
317363 ))
318-
364+ wheel_infos .sort (key = lambda info : _python_tag_sort_key (info .python_tag ))
365+
319366 if sys .platform == "linux" :
320367 combine_wheels_linux (args , wheel_infos )
321368 elif sys .platform == "win32" :
@@ -329,9 +376,10 @@ def main(args: argparse.Namespace):
329376
330377if __name__ == "__main__" :
331378 parser = argparse .ArgumentParser ()
332- parser .add_argument ("--input_folder" , type = str , required = True , help = "Path to the folder containing already repaired wheels for individual python versions that are to be combined into a single wheel" )
379+ parser .add_argument ("--input_folder" , type = str , help = "Path to the folder containing already repaired wheels for individual python versions that are to be combined into a single wheel" )
333380 parser .add_argument ("--output_folder" , type = str , default = "." , help = "Path to the folder where the combined wheel will be saved" )
334381 parser .add_argument ("--strip_unneeded" , action = "store_true" , help = "Strip the libraries to reduce size" )
335382 parser .add_argument ("--log_level" , type = str , default = "INFO" , help = "Log level" )
383+ parser .add_argument ("--self-test" , action = "store_true" , help = "Run internal tests and exit" )
336384 args = parser .parse_args ()
337- main (args )
385+ main (args )
0 commit comments