def find_usages(
package_name: str, src_dir: Path, n_processes: int, batch_size: int
) -> UsageCountStore:
python_files = list_files(src_dir, ".py")
python_file_batches = _split_into_batches(python_files, batch_size)
aggregated_counts = UsageCountStore()
for batch_index in range(0, len(python_file_batches), n_processes):
python_file_batches_slice = python_file_batches[
batch_index : batch_index + n_processes
]
n_process_to_spawn = min(n_processes, len(python_file_batches_slice))
with Pool(
processes=n_process_to_spawn,
initializer=_initializer,
initargs=[logging.root.level],
) as pool:
batch_counts = pool.starmap(
_find_usages_in_batch,
[[package_name, it] for it in python_file_batches_slice],
)
for batch_count in batch_counts:
aggregated_counts.merge_other_into_self(batch_count)
return aggregated_counts