Home > Enterprise >  How to make a loop that works with either a list or dictionary
How to make a loop that works with either a list or dictionary

Time:06-08

I have two almost identical methods, but I have no idea how to merge them into one as I don't know how to make the loop capable of looping through either a list or a dictionary.

Function 1:

def Function1()
    # a block of code common for both methods

    some_list = ['A', 'B', 'C']
    for item in some_list:
        condition = boolean_function_A(item)
        if condition:
            # a block of code common for both methods

Function2:

def Function2()
    # a block of code common for both methods

    unique_line_for_function2()
    some_dict = {'A':1, 'B':2, 'C':3}
    for key, value in some_dict.items():
        condition = boolean_function_B(key, value)
        if condition:
            # a block of code common for both methods

I would already merge them into one, but how to handle a different number of loop arguments? If I unpack a list I have one value, and if I unpack a dictionary, I have (and need) two values.

I know I can also encapsulate the common blocks of code - but still these functions do almost the same job, only the condition is different and takes different values. I would like to stick to a single method so I Don't Repeat Myself.

CodePudding user response:

Edit: Even better, you don't have to unpack the values while defining the loop at all. You could do for item in (some_collection.items() if is_dict else some_collection): and item will simply become a tuple when necessary. From there you could manually unpack item on the first line of the loop with something like if is_dict: key, value = item. Both solutions will work.


In this particular case (being list vs. dictionary), I think your best option would actually be to not use .items() on the dictionary at all, and instead do something like:

def function():
    # a block of code common to both methods

    some_collection = {'A': 1, 'B': 2, 'C': 3}
    is_dict = isinstance(some_collection, dict)

    if is_dict: unique_line_for_function2()
    else: unique_line_for_function1()

    for item in some_collection:    # for a dictionary, this will loop over the keys
        if not is_dict: condition = boolean_function_A(item)
        else: condition = boolean_function_B(item, some_collection[item])
        if condition:
            # a block of code common to both methods

A plain for-loop over a dictionary will loop over its keys (even without using .keys()), which can be used to access the values. You can expand on this however you'd like for your use-case, but this is just an idea.

CodePudding user response:

You can use tuple unpacking (*star_args) to pass a variable number of arguments to a function. So you just need to ensure that both the list and the dict supply their values as a tuple:

def function(some_data)
    # a block of code common for both methods
    
    if isinstance(some_data, dict):
        unique_line_for_function2()
        values = some_data.items()
        boolean_function = boolean_function_B
    else:
        values = ((item,) for item in some_data)
        boolean_function = boolean_function_A
    
    for args in values:
        if boolean_function(*args):
            # a block of code common for both methods
  • Related