How do I reuse the code from this function cleanly?

Question

The following function is used within a module to query network devices and is called from multiple scripts I use. The arguments it takes are a nested dictionary (the device IP and creds etc) and a string (the command to run on the device):

def query_devices(hosts, cmd):
    
    results = {
        'success' : [],
        'failed' : [],
    }

    for host in hosts:
        device = hosts[host]
        try:
            swp = ConnectHandler(device_type=device['dev_type'],
                                 ip=device['ip'],
                                 username=device['user'],
                                 password=device['pwd'],
                                 secret=device['secret'])
            swp.enable()
            results[host] = swp.send_command(cmd)
            results['success'].append(host)
            swp.disconnect()
        except (NetMikoTimeoutException, NetMikoAuthenticationException) as e:
            results['failed'].append(host)
            results[host] = e

    return results

I want to reuse all of the code to update a device and the only changes would be:

1. The function would take the same dictionary but the cmd argument would now be a list of commands.

2. The following line:

    results[host] = swp.send_command(cmd)

   would be replaced by:

    results[host] = swp.send_config_set(cmd)

I could obviously just replicate the function making those two changes and as it is in a module I reuse, I am only having to do it once but I am still basically repeating a lot of the same code.

Is there a better way to do this as I seem to come across the same issue quite often in my code.

Answer 1

You could just add a check on the changed line:

...
if isinstance(cmd, str):
    results[host] = swp.send_command(cmd)
else:
    results[host] = swp.send_config_set(cmd)
...

The rest of the function can stay the same and now you can simply call it with either a string or a list of strings...

Answer 2

You could use the unpacking operator to always make cmds an iterable (a tuple, actually) even if it is a single value. That way you could always call send_config_set . Here is a super simplified example to illustrate the concept.

def query_devices(hosts, *cmds):
    for one_cmd in cmds:
        print(one_cmd)
        
print('query_devices("hosts", "cmd_1")')
query_devices("hosts", "cmd_1")

print('\nquery_devices("hosts", "cmd_1", "cmd_2", "cmd_3")')
query_devices("hosts", "cmd_1", "cmd_2", "cmd_3")

print('\nquery_devices("hosts", *["cmd_1", "cmd_2", "cmd_3"])')
query_devices("hosts", *["cmd_1", "cmd_2", "cmd_3"])

Output:

query_devices("hosts", "cmd_1")
cmd_1

query_devices("hosts", "cmd_1", "cmd_2", "cmd_3")
cmd_1
cmd_2
cmd_3

query_devices("hosts", *["cmd_1", "cmd_2", "cmd_3"])
cmd_1
cmd_2
cmd_3