From 90ad4b30d0b3c8876fd28c64455700be49d78f71 Mon Sep 17 00:00:00 2001 From: Kyle Swanson Date: Sat, 3 Feb 2018 08:29:20 -0500 Subject: [PATCH] Allowing for single variables rather than lists and adding some tests --- .gitignore | 2 + p_tqdm/__init__.py | 206 +++++++++++++++++++-------------------- p_tqdm/tests/__init__.py | 0 p_tqdm/tests/tests.py | 158 ++++++++++++++++++++++++++++++ setup.py | 4 +- 5 files changed, 266 insertions(+), 104 deletions(-) create mode 100644 p_tqdm/tests/__init__.py create mode 100644 p_tqdm/tests/tests.py diff --git a/.gitignore b/.gitignore index b546818..5c74ea5 100644 --- a/.gitignore +++ b/.gitignore @@ -2,3 +2,5 @@ MANIFEST dist gifs .DS_Store +p_tqdm.egg-info +__pycache__ diff --git a/p_tqdm/__init__.py b/p_tqdm/__init__.py index a6de810..8732cff 100644 --- a/p_tqdm/__init__.py +++ b/p_tqdm/__init__.py @@ -12,156 +12,156 @@ from pathos.helpers import cpu_count from pathos.multiprocessing import ProcessingPool as Pool from tqdm import tqdm -def p_imap(function, *arrays, **kwargs): - """Returns an iterator for a parallel ordered map with a progress bar. +def _parallel(ordered, function, *arrays, **kwargs): + """Returns an iterator for a parallel map with a progress bar. - Args: - function: The function to apply to each element + Arguments: + ordered(bool): True for an ordered map, false for an unordered map. + function(function): The function to apply to each element of the given arrays. - arrays: One or more arrays of the same length - containing the data to be mapped. - num_cpus: The number of cpus to use in parallel. + arrays(tuple): One or more arrays of the same length + containing the data to be mapped. If a non-list + variable is passed, it will be repeated a number + of times equal to the lengths of the list(s). If only + non-list variables are passed, the function will be + performed num_iter times. + num_cpus(int): The number of cpus to use in parallel. If an int, uses that many cpus. If a float, uses that proportion of cpus. If None, uses all available cpus. + num_iter(int): If only non-list variables are passed, the + function will be performed num_iter times on + these variables. Default: 1. + Returns: An iterator which will apply the function to each element of the given arrays in parallel in order with a progress bar. """ - num_cpus = kwargs.get('num_cpus', None) + # Convert tuple to list + arrays = list(arrays) + # Extract kwargs + num_cpus = kwargs.get('num_cpus', None) + num_iter = kwargs.get('num_iter', 1) + + # Determine num_cpus if num_cpus is None: num_cpus = cpu_count() elif type(num_cpus) == float: num_cpus = int(round(num_cpus * cpu_count())) - iterator = tqdm(Pool(num_cpus).imap(function, *arrays), - total=len(arrays[0])) + # Determine num_iter when at least one list is present + if any([type(array) == list for array in arrays]): + num_iter = max([len(array) for array in arrays if type(array) == list]) + + # Convert single variables to lists + # and confirm lists are same length + for i, array in enumerate(arrays): + if type(array) != list: + arrays[i] = [array for _ in range(num_iter)] + else: + assert len(array) == num_iter + + # Create parallel iterator + map_type = 'imap' if ordered else 'uimap' + iterator = tqdm(getattr(Pool(num_cpus), map_type)(function, *arrays), + total=num_iter) + + return iterator + +def p_imap(function, *arrays, **kwargs): + """Returns an iterator for a parallel ordered map with a progress bar.""" + + ordered = True + iterator = _parallel(ordered, function, *arrays, **kwargs) return iterator def p_map(function, *arrays, **kwargs): - """Performs a parallel ordered map with a progress bar. + """Performs a parallel ordered map with a progress bar.""" - Example: - p_map(f, [1, 2, 3], ['a', 'b', 'c']) --> [f(1, 'a'), f(2, 'b'), f(3, 'c')] + ordered = True + iterator = _parallel(ordered, function, *arrays, **kwargs) + result = list(iterator) - Args: - function: The function to apply to each element - of the given arrays. - arrays: One or more arrays of the same length - containing the data to be mapped. - num_cpus: The number of cpus to use in parallel. - If an int, uses that many cpus. - If a float, uses that proportion of cpus. - If None, uses all available cpus. - Returns: - An array with the result of applying the function - to each element of the given arrays in order. - """ - - num_cpus = kwargs.get('num_cpus', None) - - new_data = list(p_imap(function, *arrays, num_cpus=num_cpus)) - - return new_data + return result def p_uimap(function, *arrays, **kwargs): - """Returns an iterator for a parallel unordered map with a progress bar. + """Returns an iterator for a parallel unordered map with a progress bar.""" - Args: - function: The function to apply to each element - of the given arrays. - arrays: One or more arrays of the same length - containing the data to be mapped. - num_cpus: The number of cpus to use in parallel. - If an int, uses that many cpus. - If a float, uses that proportion of cpus. - If None, uses all available cpus. - Returns: - An iterator which will apply the function - to each element of the given arrays in - parallel with a progress bar. The results - may be in any order. - """ - - num_cpus = kwargs.get('num_cpus', None) - - if num_cpus is None: - num_cpus = cpu_count() - elif type(num_cpus) == float: - num_cpus = int(round(num_cpus * cpu_count())) - - iterator = tqdm(Pool(num_cpus).uimap(function, *arrays), - total=len(arrays[0])) + ordered = False + iterator = _parallel(ordered, function, *arrays, **kwargs) return iterator def p_umap(function, *arrays, **kwargs): - """Performs a parallel unordered map with a progress bar. + """Performs a parallel unordered map with a progress bar.""" - Example: - p_umap(f, [1, 2, 3], ['a', 'b', 'c']) --> [f(2, 'b'), f(1, 'a'), f(3, 'c')] - Note: The resulting array may be in any order. + ordered = False + iterator = _parallel(ordered, function, *arrays, **kwargs) + result = list(iterator) - Args: - function: The function to apply to each element - of the given arrays. - arrays: One or more arrays of the same length - containing the data to be mapped. - num_cpus: The number of cpus to use in parallel. - If an int, uses that many cpus. - If a float, uses that proportion of cpus. - If None, uses all available cpus. - Returns: - An array with the result of applying the function - to each element of the given arrays. This array - may be in any order. - """ + return result - num_cpus = kwargs.get('num_cpus', None) - - new_data = list(p_uimap(function, *arrays, num_cpus=num_cpus)) - - return new_data - -def t_imap(function, *arrays): +def _sequential(function, *arrays, **kwargs): """Returns an iterator for a sequential map with a progress bar. - Args: - function: The function to apply to each element + Arguments: + function(function): The function to apply to each element of the given arrays. - arrays: One or more arrays of the same length - containing the data to be mapped. + arrays(tuple): One or more arrays of the same length + containing the data to be mapped. If a non-list + variable is passed, it will be repeated a number + of times equal to the lengths of the list(s). If only + non-list variables are passed, the function will be + performed num_iter times. + num_iter(int): If only non-list variables are passed, the + function will be performed num_iter times on + these variables. Default: 1. + Returns: An iterator which will apply the function to each element of the given arrays sequentially in order with a progress bar. """ + # Convert tuple to list + arrays = list(arrays) + + # Extract kwargs + num_iter = kwargs.get('num_iter', 1) + + # Determine num_iter when at least one list is present + if any([type(array) == list for array in arrays]): + num_iter = max([len(array) for array in arrays if type(array) == list]) + + # Convert single variables to lists + # and confirm lists are same length + for i, array in enumerate(arrays): + if type(array) != list: + arrays[i] = [array for _ in range(num_iter)] + else: + assert len(array) == num_iter + + # Create parallel iterator iterator = tqdm(map(function, *arrays), - total=len(arrays[0])) + total=num_iter) return iterator -def t_map(function, *arrays): - """Performs a sequential map with a progress bar. +def t_imap(function, *arrays, **kwargs): + """Returns an iterator for a sequential map with a progress bar.""" - Example: - t_map(f, [1, 2, 3], ['a', 'b', 'c']) --> [f(1, 'a'), f(2, 'b'), f(3, 'c')] + iterator = sequential(function, *arrays, **kwargs) - Args: - function: The function to apply to each element - of the given arrays. - arrays: One or more arrays of the same length - containing the data to be mapped. - Returns: - An array with the result of applying the function - to each element of the given arrays in order. - """ + return iterator - new_data = list(t_imap(function, *arrays)) +def t_map(function, *arrays, **kwargs): + """Performs a sequential map with a progress bar.""" - return new_data + iterator = sequential(function, *arrays, **kwargs) + result = list(iteratorZ) + + return result diff --git a/p_tqdm/tests/__init__.py b/p_tqdm/tests/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/p_tqdm/tests/tests.py b/p_tqdm/tests/tests.py new file mode 100644 index 0000000..2539ba0 --- /dev/null +++ b/p_tqdm/tests/tests.py @@ -0,0 +1,158 @@ +import types +import unittest + +import p_tqdm +import tqdm + +def add_1(a): + return a + 1 + +def add_2(a, b): + return a + b + +def add_3(a, b, c): + return a + b + c + +def _test_one_list(self): + array = [1, 2, 3] + result = self.func(add_1, array) + if self.generator: + result = list(result) + + correct_array = [2, 3, 4] + self.assertEqual(correct_array, result) + +def _test_two_lists(self): + array_1 = [1, 2, 3] + array_2 = [10, 11, 12] + result = self.func(add_2, array_1, array_2) + if self.generator: + result = list(result) + + correct_array = [11, 13, 15] + self.assertEqual(correct_array, result) + +def _test_two_lists_and_one_single(self): + array_1 = [1, 2, 3] + array_2 = [10, 11, 12] + single = 5 + result = self.func(add_3, array_1, single, array_2) + if self.generator: + result = list(result) + + correct_array = [16, 18, 20] + self.assertEqual(correct_array, result) + +def _test_one_list_and_two_singles(self): + array = [1, 2, 3] + single_1 = 5 + single_2 = -2 + result = self.func(add_3, single_1, array, single_2) + if self.generator: + result = list(result) + + correct_array = [4, 5, 6] + self.assertEqual(correct_array, result) + +def _test_one_single(self): + single = 5 + result = self.func(add_1, single) + if self.generator: + result = list(result) + + correct_array = [6] + self.assertEqual(correct_array, result) + +def _test_one_single_with_num_iter(self): + single = 5 + num_iter = 3 + result = self.func(add_1, single, num_iter=num_iter) + if self.generator: + result = list(result) + + correct_array = [6]*num_iter + self.assertEqual(correct_array, result) + +def _test_two_singles(self): + single_1 = 5 + single_2 = -2 + result = self.func(add_2, single_1, single_2) + if self.generator: + result = list(result) + + correct_array = [3] + self.assertEqual(correct_array, result) + +def _test_two_singles_with_num_iter(self): + single_1 = 5 + single_2 = -2 + num_iter = 3 + result = self.func(add_2, single_1, single_2, num_iter=num_iter) + if self.generator: + result = list(result) + + correct_array = [3]*num_iter + self.assertEqual(correct_array, result) + +class Testp_imap(unittest.TestCase): + def __init__(self, *args, **kwargs): + super(Testp_imap, self).__init__(*args, **kwargs) + self.func = p_tqdm.p_imap + self.generator = True + + def test_one_list(self): + _test_one_list(self) + + def test_two_lists(self): + _test_two_lists(self) + + def test_two_lists_and_one_single(self): + _test_two_lists_and_one_single(self) + + def test_one_list_and_two_singles(self): + _test_one_list_and_two_singles(self) + + def test_one_single(self): + _test_one_single(self) + + def test_one_single_with_num_iter(self): + _test_one_single_with_num_iter(self) + + def test_two_singles(self): + _test_two_singles(self) + + def test_two_singles_with_num_iter(self): + _test_two_singles_with_num_iter(self) + +class Testp_map(unittest.TestCase): + def __init__(self, *args, **kwargs): + super(Testp_map, self).__init__(*args, **kwargs) + self.func = p_tqdm.p_map + self.generator = False + + def test_one_list(self): + _test_one_list(self) + + def test_two_lists(self): + _test_two_lists(self) + + def test_two_lists_and_one_single(self): + _test_two_lists_and_one_single(self) + + def test_one_list_and_two_singles(self): + _test_one_list_and_two_singles(self) + + def test_one_single(self): + _test_one_single(self) + + def test_one_single_with_num_iter(self): + _test_one_single_with_num_iter(self) + + def test_two_singles(self): + _test_two_singles(self) + + def test_two_singles_with_num_iter(self): + _test_two_singles_with_num_iter(self) + +if __name__ == '__main__': + unittest.main() diff --git a/setup.py b/setup.py index 730d75b..86d3ec4 100644 --- a/setup.py +++ b/setup.py @@ -1,4 +1,4 @@ -from distutils.core import setup +from setuptools import setup setup( name = 'p_tqdm', @@ -10,6 +10,8 @@ setup( url = 'https://github.com/swansonk14/p_tqdm', license = 'MIT', install_requires = ['tqdm', 'pathos'], + test_suite='nose.collector', + tests_require=['nose'], keywords = ['tqdm', 'progress bar', 'parallel'], classifiers = [], )