# Path: monty/multiprocessing.py
# def imap_tqdm(nprocs: int, func: Callable, iterable: Iterable, *args, **kwargs):
#     """
#     A wrapper around Pool.imap. Creates a Pool with nprocs and then runs a f
#     unction over an iterable with progress bar.
# 
#     :param nprocs: Number of processes
#     :param func: Callable
#     :param iterable: Iterable of arguments.
#     :param args: Passthrough to Pool.imap
#     :param kwargs: Passthrough to Pool.imap
#     :return: Results of Pool.imap.
#     """
#     if tqdm is None:
#         raise ImportError("tqdm must be installed for this function.")
#     data = []
#     with Pool(nprocs) as pool:
#         try:
#             n = len(iterable)  # type: ignore
#         except TypeError:
#             n = None  # type: ignore
#         with tqdm(total=n) as pbar:
#             for i, d in enumerate(pool.imap(func, iterable, *args, **kwargs)):
#                 pbar.update()
#                 data.append(d)
#     return data


import unittest
from monty.multiprocessing import imap_tqdm
from math import sqrt



class FuncCase(unittest.TestCase):
    def test_imap_tqdm(self):
        results = imap_tqdm(4, sqrt, range(10000))
        self.assertEqual(len(results), 10000)
        self.assertEqual(results[0], 0)
        self.assertEqual(results[400], 20)
        self.assertEqual(results[9999], 99.99499987499375)

# Path: monty/multiprocessing.py
# def imap_tqdm(nprocs: int, func: Callable, iterable: Iterable, *args, **kwargs):
#     """
#     A wrapper around Pool.imap. Creates a Pool with nprocs and then runs a f
#     unction over an iterable with progress bar.
# 
#     :param nprocs: Number of processes
#     :param func: Callable
#     :param iterable: Iterable of arguments.
#     :param args: Passthrough to Pool.imap
#     :param kwargs: Passthrough to Pool.imap
#     :return: Results of Pool.imap.
#     """
#     if tqdm is None:
#         raise ImportError("tqdm must be installed for this function.")
#     data = []
#     with Pool(nprocs) as pool:
#         try:
#             n = len(iterable)  # type: ignore
#         except TypeError:
#             n = None  # type: ignore
#         with tqdm(total=n) as pbar:
#             for i, d in enumerate(pool.imap(func, iterable, *args, **kwargs)):
#                 pbar.update()
#                 data.append(d)
#     return data


# Path: tests/test_multiprocessing.py
import unittest
from monty.multiprocessing import imap_tqdm
from math import sqrt




class FuncCase(unittest.TestCase):
    def test_imap_tqdm(self):
        results = imap_tqdm(4, sqrt, range(10000))
        self.assertEqual(len(results), 10000)
        self.assertEqual(results[0], 0)
        self.assertEqual(results[400], 20)
        self.assertEqual(results[9999], 99.99499987499375)

        results = imap_tqdm(4, sqrt, (i**2 for i in range(10000)))