4
4
from __future__ import annotations
5
5
6
6
import argparse
7
+ import concurrent .futures
7
8
import os
8
9
import re
10
+ import subprocess
9
11
import sys
10
12
import tempfile
11
- from contextlib import redirect_stderr , redirect_stdout
13
+ import time
14
+ from collections import defaultdict
12
15
from dataclasses import dataclass
13
- from io import StringIO
14
16
from itertools import product
15
17
from pathlib import Path
18
+ from threading import Lock
16
19
from typing import TYPE_CHECKING , Any , NamedTuple
17
20
18
21
if TYPE_CHECKING :
23
26
import tomli
24
27
from utils import (
25
28
VERSIONS_RE as VERSION_LINE_RE ,
29
+ PackageDependencies ,
30
+ VenvInfo ,
26
31
colored ,
27
32
get_gitignore_spec ,
33
+ get_mypy_req ,
28
34
get_recursive_requirements ,
35
+ make_venv ,
29
36
print_error ,
30
37
print_success_msg ,
31
38
spec_matches_path ,
32
39
strip_comments ,
33
40
)
34
41
42
+ # Fail early if mypy isn't installed
35
43
try :
36
- from mypy . api import run as mypy_run
44
+ import mypy # noqa: F401
37
45
except ImportError :
38
46
print_error ("Cannot import mypy. Did you install it?" )
39
47
sys .exit (1 )
@@ -108,7 +116,7 @@ class TestConfig:
108
116
109
117
def log (args : TestConfig , * varargs : object ) -> None :
110
118
if args .verbose >= 2 :
111
- print (* varargs )
119
+ print (colored ( " " . join ( map ( str , varargs )), "blue" ) )
112
120
113
121
114
122
def match (path : Path , args : TestConfig ) -> bool :
@@ -204,7 +212,19 @@ def add_configuration(configurations: list[MypyDistConf], distribution: str) ->
204
212
configurations .append (MypyDistConf (module_name , values .copy ()))
205
213
206
214
207
- def run_mypy (args : TestConfig , configurations : list [MypyDistConf ], files : list [Path ], * , testing_stdlib : bool ) -> ReturnCode :
215
+ def run_mypy (
216
+ args : TestConfig ,
217
+ configurations : list [MypyDistConf ],
218
+ files : list [Path ],
219
+ * ,
220
+ testing_stdlib : bool ,
221
+ non_types_dependencies : bool ,
222
+ venv_info : VenvInfo ,
223
+ mypypath : str | None = None ,
224
+ ) -> ReturnCode :
225
+ env_vars = dict (os .environ )
226
+ if mypypath is not None :
227
+ env_vars ["MYPYPATH" ] = mypypath
208
228
with tempfile .NamedTemporaryFile ("w+" ) as temp :
209
229
temp .write ("[mypy]\n " )
210
230
for dist_conf in configurations :
@@ -213,57 +233,49 @@ def run_mypy(args: TestConfig, configurations: list[MypyDistConf], files: list[P
213
233
temp .write (f"{ k } = { v } \n " )
214
234
temp .flush ()
215
235
216
- flags = get_mypy_flags (args , temp .name , testing_stdlib = testing_stdlib )
236
+ flags = [
237
+ "--python-version" ,
238
+ args .version ,
239
+ "--show-traceback" ,
240
+ "--warn-incomplete-stub" ,
241
+ "--no-error-summary" ,
242
+ "--platform" ,
243
+ args .platform ,
244
+ "--custom-typeshed-dir" ,
245
+ str (Path (__file__ ).parent .parent ),
246
+ "--strict" ,
247
+ # Stub completion is checked by pyright (--allow-*-defs)
248
+ "--allow-untyped-defs" ,
249
+ "--allow-incomplete-defs" ,
250
+ "--allow-subclassing-any" , # TODO: Do we still need this now that non-types dependencies are allowed? (#5768)
251
+ "--enable-error-code" ,
252
+ "ignore-without-code" ,
253
+ "--config-file" ,
254
+ temp .name ,
255
+ ]
256
+ if not testing_stdlib :
257
+ flags .append ("--explicit-package-bases" )
258
+ if not non_types_dependencies :
259
+ flags .append ("--no-site-packages" )
260
+
217
261
mypy_args = [* flags , * map (str , files )]
262
+ mypy_command = [venv_info .python_exe , "-m" , "mypy" ] + mypy_args
218
263
if args .verbose :
219
- print ("running mypy" , " " .join (mypy_args ))
220
- stdout_redirect , stderr_redirect = StringIO (), StringIO ()
221
- with redirect_stdout (stdout_redirect ), redirect_stderr (stderr_redirect ):
222
- returned_stdout , returned_stderr , exit_code = mypy_run (mypy_args )
223
-
224
- if exit_code :
264
+ print (colored (f"running { ' ' .join (mypy_command )} " , "blue" ))
265
+ result = subprocess .run (mypy_command , capture_output = True , text = True , env = env_vars )
266
+ if result .returncode :
225
267
print_error ("failure\n " )
226
- captured_stdout = stdout_redirect .getvalue ()
227
- captured_stderr = stderr_redirect .getvalue ()
228
- if returned_stderr :
229
- print_error (returned_stderr )
230
- if captured_stderr :
231
- print_error (captured_stderr )
232
- if returned_stdout :
233
- print_error (returned_stdout )
234
- if captured_stdout :
235
- print_error (captured_stdout , end = "" )
268
+ if result .stdout :
269
+ print_error (result .stdout )
270
+ if result .stderr :
271
+ print_error (result .stderr )
272
+ if non_types_dependencies and args .verbose :
273
+ print ("Ran with the following environment:" )
274
+ subprocess .run ([venv_info .pip_exe , "freeze" , "--all" ])
275
+ print ()
236
276
else :
237
277
print_success_msg ()
238
- return exit_code
239
-
240
-
241
- def get_mypy_flags (args : TestConfig , temp_name : str , * , testing_stdlib : bool ) -> list [str ]:
242
- flags = [
243
- "--python-version" ,
244
- args .version ,
245
- "--show-traceback" ,
246
- "--warn-incomplete-stub" ,
247
- "--show-error-codes" ,
248
- "--no-error-summary" ,
249
- "--platform" ,
250
- args .platform ,
251
- "--no-site-packages" ,
252
- "--custom-typeshed-dir" ,
253
- str (Path (__file__ ).parent .parent ),
254
- "--strict" ,
255
- # Stub completion is checked by pyright (--allow-*-defs)
256
- "--allow-untyped-defs" ,
257
- "--allow-incomplete-defs" ,
258
- "--allow-subclassing-any" , # Needed until we can use non-types dependencies #5768
259
- "--enable-error-code" ,
260
- "ignore-without-code" ,
261
- "--config-file" ,
262
- temp_name ,
263
- ]
264
- if not testing_stdlib :
265
- flags .append ("--explicit-package-bases" )
266
- return flags
278
+ return result .returncode
267
279
268
280
269
281
def add_third_party_files (
@@ -298,7 +310,9 @@ class TestResults(NamedTuple):
298
310
files_checked : int
299
311
300
312
301
- def test_third_party_distribution (distribution : str , args : TestConfig ) -> TestResults :
313
+ def test_third_party_distribution (
314
+ distribution : str , args : TestConfig , venv_info : VenvInfo , * , non_types_dependencies : bool
315
+ ) -> TestResults :
302
316
"""Test the stubs of a third-party distribution.
303
317
304
318
Return a tuple, where the first element indicates mypy's return code
@@ -313,20 +327,24 @@ def test_third_party_distribution(distribution: str, args: TestConfig) -> TestRe
313
327
if not files and args .filter :
314
328
return TestResults (0 , 0 )
315
329
316
- print (f"testing { distribution } ({ len (files )} files)... " , end = "" )
330
+ print (f"testing { distribution } ({ len (files )} files)... " , end = "" , flush = True )
317
331
318
332
if not files :
319
333
print_error ("no files found" )
320
334
sys .exit (1 )
321
335
322
- prev_mypypath = os .getenv ("MYPYPATH" )
323
- os .environ ["MYPYPATH" ] = os .pathsep .join (str (Path ("stubs" , dist )) for dist in seen_dists )
324
- code = run_mypy (args , configurations , files , testing_stdlib = False )
325
- if prev_mypypath is None :
326
- del os .environ ["MYPYPATH" ]
327
- else :
328
- os .environ ["MYPYPATH" ] = prev_mypypath
329
-
336
+ mypypath = os .pathsep .join (str (Path ("stubs" , dist )) for dist in seen_dists )
337
+ if args .verbose :
338
+ print (colored (f"\n { mypypath = } " , "blue" ))
339
+ code = run_mypy (
340
+ args ,
341
+ configurations ,
342
+ files ,
343
+ venv_info = venv_info ,
344
+ mypypath = mypypath ,
345
+ testing_stdlib = False ,
346
+ non_types_dependencies = non_types_dependencies ,
347
+ )
330
348
return TestResults (code , len (files ))
331
349
332
350
@@ -343,19 +361,105 @@ def test_stdlib(code: int, args: TestConfig) -> TestResults:
343
361
add_files (files , (stdlib / name ), args )
344
362
345
363
if files :
346
- print (f"Testing stdlib ({ len (files )} files)..." )
347
- print ("Running mypy " + " " .join (get_mypy_flags (args , "/tmp/..." , testing_stdlib = True )))
348
- this_code = run_mypy (args , [], files , testing_stdlib = True )
364
+ print (f"Testing stdlib ({ len (files )} files)..." , end = "" , flush = True )
365
+ # We don't actually need pip for the stdlib testing
366
+ venv_info = VenvInfo (pip_exe = "" , python_exe = sys .executable )
367
+ this_code = run_mypy (args , [], files , venv_info = venv_info , testing_stdlib = True , non_types_dependencies = False )
349
368
code = max (code , this_code )
350
369
351
370
return TestResults (code , len (files ))
352
371
353
372
354
- def test_third_party_stubs (code : int , args : TestConfig ) -> TestResults :
373
+ _PRINT_LOCK = Lock ()
374
+ _DISTRIBUTION_TO_VENV_MAPPING : dict [str , VenvInfo ] = {}
375
+
376
+
377
+ def setup_venv_for_external_requirements_set (requirements_set : frozenset [str ], tempdir : Path ) -> tuple [frozenset [str ], VenvInfo ]:
378
+ venv_dir = tempdir / f".venv-{ hash (requirements_set )} "
379
+ return requirements_set , make_venv (venv_dir )
380
+
381
+
382
+ def install_requirements_for_venv (venv_info : VenvInfo , args : TestConfig , external_requirements : frozenset [str ]) -> None :
383
+ # Use --no-cache-dir to avoid issues with concurrent read/writes to the cache
384
+ pip_command = [venv_info .pip_exe , "install" , get_mypy_req (), * sorted (external_requirements ), "--no-cache-dir" ]
385
+ if args .verbose :
386
+ with _PRINT_LOCK :
387
+ print (colored (f"Running { pip_command } " , "blue" ))
388
+ try :
389
+ subprocess .run (pip_command , check = True , capture_output = True , text = True )
390
+ except subprocess .CalledProcessError as e :
391
+ print (e .stderr )
392
+ raise
393
+
394
+
395
+ def setup_virtual_environments (distributions : dict [str , PackageDependencies ], args : TestConfig , tempdir : Path ) -> None :
396
+ # We don't actually need pip if there aren't any external dependencies
397
+ no_external_dependencies_venv = VenvInfo (pip_exe = "" , python_exe = sys .executable )
398
+ external_requirements_to_distributions : defaultdict [frozenset [str ], list [str ]] = defaultdict (list )
399
+ num_pkgs_with_external_reqs = 0
400
+
401
+ for distribution_name , requirements in distributions .items ():
402
+ if requirements .external_pkgs :
403
+ num_pkgs_with_external_reqs += 1
404
+ external_requirements = frozenset (requirements .external_pkgs )
405
+ external_requirements_to_distributions [external_requirements ].append (distribution_name )
406
+ else :
407
+ _DISTRIBUTION_TO_VENV_MAPPING [distribution_name ] = no_external_dependencies_venv
408
+
409
+ if num_pkgs_with_external_reqs == 0 :
410
+ if args .verbose :
411
+ print (colored ("No additional venvs are required to be set up" , "blue" ))
412
+ return
413
+
414
+ requirements_sets_to_venvs : dict [frozenset [str ], VenvInfo ] = {}
415
+
416
+ if args .verbose :
417
+ num_venvs = len (external_requirements_to_distributions )
418
+ msg = (
419
+ f"Setting up { num_venvs } venv{ 's' if num_venvs != 1 else '' } "
420
+ f"for { num_pkgs_with_external_reqs } "
421
+ f"distribution{ 's' if num_pkgs_with_external_reqs != 1 else '' } ... "
422
+ )
423
+ print (colored (msg , "blue" ), end = "" , flush = True )
424
+ venv_start_time = time .perf_counter ()
425
+
426
+ with concurrent .futures .ThreadPoolExecutor () as executor :
427
+ venv_info_futures = [
428
+ executor .submit (setup_venv_for_external_requirements_set , requirements_set , tempdir )
429
+ for requirements_set in external_requirements_to_distributions
430
+ ]
431
+ for venv_info_future in concurrent .futures .as_completed (venv_info_futures ):
432
+ requirements_set , venv_info = venv_info_future .result ()
433
+ requirements_sets_to_venvs [requirements_set ] = venv_info
434
+
435
+ if args .verbose :
436
+ venv_elapsed_time = time .perf_counter () - venv_start_time
437
+ print (colored (f"took { venv_elapsed_time :.2f} seconds" , "blue" ))
438
+ pip_start_time = time .perf_counter ()
439
+
440
+ # Limit workers to 10 at a time, since this makes network requests
441
+ with concurrent .futures .ThreadPoolExecutor (max_workers = 10 ) as executor :
442
+ pip_install_futures = [
443
+ executor .submit (install_requirements_for_venv , venv_info , args , requirements_set )
444
+ for requirements_set , venv_info in requirements_sets_to_venvs .items ()
445
+ ]
446
+ concurrent .futures .wait (pip_install_futures )
447
+
448
+ if args .verbose :
449
+ pip_elapsed_time = time .perf_counter () - pip_start_time
450
+ msg = f"Combined time for installing requirements across all venvs: { pip_elapsed_time :.2f} seconds"
451
+ print (colored (msg , "blue" ))
452
+
453
+ for requirements_set , distribution_list in external_requirements_to_distributions .items ():
454
+ venv_to_use = requirements_sets_to_venvs [requirements_set ]
455
+ _DISTRIBUTION_TO_VENV_MAPPING .update (dict .fromkeys (distribution_list , venv_to_use ))
456
+
457
+
458
+ def test_third_party_stubs (code : int , args : TestConfig , tempdir : Path ) -> TestResults :
355
459
print ("Testing third-party packages..." )
356
- print ("Running mypy " + " " .join (get_mypy_flags (args , "/tmp/..." , testing_stdlib = False )))
357
460
files_checked = 0
358
461
gitignore_spec = get_gitignore_spec ()
462
+ distributions_to_check : dict [str , PackageDependencies ] = {}
359
463
360
464
for distribution in sorted (os .listdir ("stubs" )):
361
465
distribution_path = Path ("stubs" , distribution )
@@ -368,14 +472,25 @@ def test_third_party_stubs(code: int, args: TestConfig) -> TestResults:
368
472
or Path ("stubs" ) in args .filter
369
473
or any (distribution_path in path .parents for path in args .filter )
370
474
):
371
- this_code , checked = test_third_party_distribution (distribution , args )
372
- code = max (code , this_code )
373
- files_checked += checked
475
+ distributions_to_check [distribution ] = get_recursive_requirements (distribution )
476
+
477
+ if not _DISTRIBUTION_TO_VENV_MAPPING :
478
+ setup_virtual_environments (distributions_to_check , args , tempdir )
479
+
480
+ assert _DISTRIBUTION_TO_VENV_MAPPING .keys () == distributions_to_check .keys ()
481
+
482
+ for distribution , venv_info in _DISTRIBUTION_TO_VENV_MAPPING .items ():
483
+ non_types_dependencies = venv_info .python_exe != sys .executable
484
+ this_code , checked = test_third_party_distribution (
485
+ distribution , args , venv_info = venv_info , non_types_dependencies = non_types_dependencies
486
+ )
487
+ code = max (code , this_code )
488
+ files_checked += checked
374
489
375
490
return TestResults (code , files_checked )
376
491
377
492
378
- def test_typeshed (code : int , args : TestConfig ) -> TestResults :
493
+ def test_typeshed (code : int , args : TestConfig , tempdir : Path ) -> TestResults :
379
494
print (f"*** Testing Python { args .version } on { args .platform } " )
380
495
files_checked_this_version = 0
381
496
stdlib_dir , stubs_dir = Path ("stdlib" ), Path ("stubs" )
@@ -385,7 +500,7 @@ def test_typeshed(code: int, args: TestConfig) -> TestResults:
385
500
print ()
386
501
387
502
if stubs_dir in args .filter or any (stubs_dir in path .parents for path in args .filter ):
388
- code , third_party_files_checked = test_third_party_stubs (code , args )
503
+ code , third_party_files_checked = test_third_party_stubs (code , args , tempdir )
389
504
files_checked_this_version += third_party_files_checked
390
505
print ()
391
506
@@ -400,10 +515,12 @@ def main() -> None:
400
515
exclude = args .exclude or []
401
516
code = 0
402
517
total_files_checked = 0
403
- for version , platform in product (versions , platforms ):
404
- config = TestConfig (args .verbose , filter , exclude , version , platform )
405
- code , files_checked_this_version = test_typeshed (code , args = config )
406
- total_files_checked += files_checked_this_version
518
+ with tempfile .TemporaryDirectory () as td :
519
+ td_path = Path (td )
520
+ for version , platform in product (versions , platforms ):
521
+ config = TestConfig (args .verbose , filter , exclude , version , platform )
522
+ code , files_checked_this_version = test_typeshed (code , args = config , tempdir = td_path )
523
+ total_files_checked += files_checked_this_version
407
524
if code :
408
525
print_error (f"--- exit status { code } , { total_files_checked } files checked ---" )
409
526
sys .exit (code )
@@ -417,5 +534,5 @@ def main() -> None:
417
534
try :
418
535
main ()
419
536
except KeyboardInterrupt :
420
- print_error ("\n \n !!! \ n Test aborted due to KeyboardInterrupt\n !! !" )
537
+ print_error ("\n \n Test aborted due to KeyboardInterrupt!" )
421
538
sys .exit (1 )
0 commit comments