Nested loops, we all love them. Or do we? Five levels of indentation later you are not so sure. Your code starts to look ugly and you wish there was another way. Fortunately there is! One of the functions provided by itertools, product(), can replace your loops with a function call.
So what is itertools.product? It's a function that takes a number of iterables and returns their Cartesian product, or in simpler terms, all ordered tuples with elements coming from each of the iterables. Product of [1, 2] and ['a', 'b', 'c'] would result in tuples: (1, 'a'), (1, 'b'), (1, 'c'), (2, 'a'), (2, 'b'), (2, 'c').
It's worth mentioning that itertools.product is optimised for speed and memory usage, as the values are only computed when they're accessed. So not only will you make your program look better, but you might also make it faster and more memory efficient.
Now let's see some examples.
Example 1
In the first example I make a function call to retrieve prefixes for all combinations of regions and roles.
The benefits of using itertools.product are not immediately obvious here, but I think that it makes intent more clear. We group dependent variables together and can access values from iterables in one go. It would also be trivial to add another argument in the future if required.
pfx_list = []
for region in regions:
for role in pfx_roles:
pfx_list.extend(get_pfxs(region=region, role=role))
pfx_list = []
for region, role in itertools.product(regions, pfx_roles):
pfx_list.extend(get_pfxs(region=region, role=role))
Full code, including mock of the function, is included at the end of this post.
Example 2
Another example deals with generation of the network device names that are built from various components.
Here we are reducing four levels of nesting to one and it is immediately apparent that itertools.product version looks better, making it easier to grok what's going on.
I also did not originally iterate over range(1, 3), and it was much easier to add it to the itertools.product version once I decided to add id to device names.
iata_cc = ['bai', 'mex', 'tyo', ]
dev_types = ['switch', 'router']
dev_roles = ['core', 'edge']
for cc in iata_cc:
for dev_t in dev_types:
for dev_r in dev_roles:
for dev_id in range(1, 3):
print('{cc}-{dev_t}-{dev_r}-{dev_id}'.
format(cc=cc, dev_t=dev_t, dev_r=dev_r, dev_id=dev_id))
for cc, dev_t, dev_r, dev_id in itertools.product(iata_cc, dev_types, dev_roles, range(1, 3)):
print('{cc}-{dev_t}-{dev_r}-{dev_id}'.
format(cc=cc, dev_t=dev_t, dev_r=dev_r, dev_id=dev_id))
Exmaple 1 - full source code
def get_pfxs(region, role):
prefixes = {
'as': {
'data': ['10.1.2.0/25', '10.1.3.0/25'],
'voice': ['10.1.2.128/25', '10.1.3.128/25'],
},
'eu': {
'data': ['10.2.6.0/25', '10.2.7.0/25'],
'voice': ['10.2.6.128/25', '10.2.7.128/25'],
},
'na': {
'data': ['10.3.11.0/25', '10.3.11.0/25'],
'voice': ['10.3.12.128/25', '10.3.12.128/25'],
}
}
return prefixes[region][role]
regions = ['as', 'eu', 'na']
pfx_roles = ['data', 'voice']
pfx_list = []
for region in regions:
for role in pfx_roles:
pfx_list.extend(get_pfxs(region=region, role=role))
pfx_list = []
for region, role in itertools.product(regions, pfx_roles):
pfx_list.extend(get_pfxs(region=region, role=role))
You can get source code for this post, and an accompanying Jupyter notebook, from my GitHub repository: https://github.com/progala/pytips/tree/master/PyTips_2_itertools.product