Parallelise import scanning

Author: Peter554

src/grimp/application/usecases.py

@@ -1,7 +1,11 @@
 """
 Use cases handle application logic.
 """
-from typing import Dict, Sequence, Set, Type, Union, cast
+from typing import Dict, Sequence, Set, Type, Union, cast, List
+import multiprocessing
+import math
+
+from joblib import Parallel, delayed, parallel_config  # type: ignore
 
 from ..application.ports import caching
 from ..application.ports.filesystem import AbstractFileSystem
@@ -13,6 +17,9 @@
 from .config import settings
 
 
+N_CPUS = multiprocessing.cpu_count()
+
+
 class NotSupplied:
     pass
 
@@ -137,11 +144,23 @@ def _scan_packages(
             else:
                 imports_by_module_file[module_file] = direct_imports
 
-    # TODO Parallelise this part.
-    for module_file in modules_files_to_scan.difference(imports_by_module_file):
-        imports_by_module_file[module_file] = import_scanner.scan_for_imports(
-            module_file.module, exclude_type_checking_imports=exclude_type_checking_imports
+    remaining_modules_files_to_scan = list(
+        modules_files_to_scan.difference(imports_by_module_file)
+    )
+    chunked_remaining_modules_files_to_scan = _create_chunks(
+        remaining_modules_files_to_scan, n_chunks=N_CPUS
+    )
+    with parallel_config(n_jobs=N_CPUS):
+        import_scanning_jobs = Parallel()(
+            delayed(_scan_imports)(
+                import_scanner=import_scanner,
+                exclude_type_checking_imports=exclude_type_checking_imports,
+                module_files=chunk_module_files_to_scan,
+            )
+            for chunk_module_files_to_scan in chunked_remaining_modules_files_to_scan
         )
+        for chunk_imports_by_module_file in import_scanning_jobs:
+            imports_by_module_file.update(chunk_imports_by_module_file)
 
     imports_by_module: Dict[Module, Set[DirectImport]] = {
         k.module: v for k, v in imports_by_module_file.items()
@@ -184,3 +203,21 @@ def _is_external(module: Module, found_packages: Set[FoundPackage]) -> bool:
         module.is_descendant_of(package_module) or module == package_module
         for package_module in package_modules
     )
+
+
+def _create_chunks(module_files: List[ModuleFile], *, n_chunks: int) -> List[List[ModuleFile]]:
+    chunk_size = math.ceil(len(module_files) / n_chunks)
+    return [module_files[i * chunk_size : (i + 1) * chunk_size] for i in range(n_chunks)]
+
+
+def _scan_imports(
+    import_scanner: AbstractImportScanner,
+    exclude_type_checking_imports: bool,
+    module_files: List[ModuleFile],
+) -> Dict[ModuleFile, Set[DirectImport]]:
+    return {
+        module_file: import_scanner.scan_for_imports(
+            module_file.module, exclude_type_checking_imports=exclude_type_checking_imports
+        )
+        for module_file in module_files
+    }