Search code examples
pythondesign-patternsrefactoring

How can I avoid repeating a for loop in two different functions


I am writing an ImageCollection class in python that should hold a dictionary with a name and the image-object (pygame.image object).

In one case I want to load all images inside a folder to the dictionary and in another case just specific files, for example only button-files.

What I have written so far is this:

class ImageCollection:
    def __init__(self):
        self.dict = {}
    
    def load_images(self, path):
        directory = os.fsencode(path)

        for file in os.listdir(directory):
            file_name = os.fsdecode(file)
            img_path = path + "/" + file_name

            if file_name.endswith(".jpg") or file_name.endswith(".png"):
                # Remove extension for dictionary entry name and add image to dictionary
                #-----------------------------------------------------------------------
                dict_entry_name = file_name.removesuffix(".jpg").removesuffix(".png")
                self.dict.update({dict_entry_name: image.Image(img_path, 0)})
    
    def load_specific_images(self, path, contains_str):
        directory = os.fsencode(path)

        for file in os.listdir(directory):
            file_name = os.fsdecode(file)
            img_path = path + "/" + file_name

            if file_name.endswith(".jpg") or file_name.endswith(".png"):
                if file_name.rfind(contains_str):
                    # Remove extension for dictionary entry name and add image to dictionary
                    #-----------------------------------------------------------------------
                    dict_entry_name = file_name.removesuffix(".jpg").removesuffix(".png")
                    self.dict.update({dict_entry_name: image.Image(img_path, 0)})

The only problem is that this is probably bad programming pattern, right? In this case it probably doesnt matter but I would like to know what the best-practice in this case would be.

How can I avoid repeating myself in two different functions when the only difference is just a single if condition?

I have tried creating a "dict_add" function that creates the entry. Then I was thinking I could create two different functions, one which directly calls "dict_add" and the other one checks for the specific condition and then calls "dict_add". Then I thought I could add create just a single function with the for-loop but pass a function as an argument (which would be a callback I assume?). But one callback would need an additional argument so thats where I got stuck and wondered if my approach was correct.


Solution

  • You could make the contains_str an optional argument.

    • In cases where you want to load_images - you just provide the path
    • In cases where you want to load specific images - you provide the path and the contains_str argument

    In both cases you call load_images(...)

    Code:

    class ImageCollection:
        def __init__(self):
            self.dict = {}
        
        def load_images(self, path, contains_str=""):
            directory = os.fsencode(path)
    
            for file in os.listdir(directory):
                file_name = os.fsdecode(file)
                img_path = path + "/" + file_name
    
                if file_name.endswith(".jpg") or file_name.endswith(".png"):
                    if contains_str == "" or (contains_str != "" and file_name.rfind(contains_str)):
                        # Remove extension for dictionary entry name and add image to dictionary
                        #-----------------------------------------------------------------------
                        dict_entry_name = file_name.removesuffix(".jpg").removesuffix(".png")
                        self.dict.update({dict_entry_name: image.Image(img_path, 0)})