How to find all files with specific file extensions in a directory in Python

Question

Let's say I have a folder with several files in it with different file types: .jpg, .jpeg, .png, .tif, .gif, .docx, .pptx, .xlsx, .mp4, .avi, .mpeg, etc.

How should I set up a function that searches for a specific set of file types in a specified directory?

Let's say I only want the image files (ie jpg, jpeg, png, tif, and gif). Or I only want the video files (.mp4, .avi, .mpeg).

Would I have to write a separate function for each of these file types? Or can I have a single function that searches for image files, video files, etc.?

import os

def enterFilePath():
    global filepath
    filepath = input("Please enter your file path. ")

enterFilePath()

def enterFileName():
    global name
    name = input("Name the file. "))

enterFileName()

def data_list():
    for elem in os.listdir(filepath):
        if elem.endswith('.jpg'):
            listItem = elem + '\n'
            listName = filepath + (r"\{}List.txt".format(name))
            writeFile = open(listName, 'a')
            writeFile.write(listItem)
            writeFile.close()
        if elem.endswith('.jpeg'):
            listItem = elem + '\n'
            listName = filepath + (r"\{}List.txt".format(name))
            writeFile = open(listName, 'a')
            writeFile.write(listItem)
            writeFile.close()
        if elem.endswith('.png'):
            listItem = elem + '\n'
            listName = filepath + (r"\{}List.txt".format(name))
            writeFile = open(listName, 'a')
            writeFile.write(listItem)
            writeFile.close()
        if elem.endswith('.tif'):
            listItem = elem + '\n'
            listName = filepath + (r"\{}List.txt".format(name))
            writeFile = open(listName, 'a')
            writeFile.write(listItem)
            writeFile.close()
        if elem.endswith('.gif'):
            listItem = elem + '\n'
            listName = filepath + (r"\{}List.txt".format(name))
            writeFile = open(listName, 'a')
            writeFile.write(listItem)
            writeFile.close()
        else:
            continue
data_list()

Answer 1

Collect all your extensions in a dict and then you can just pull them out quickly and easily. I am using os.cwd() but you can just replace that with your directory.

import os

search = {}

for f in os.listdir(os.getcwd()):
    fn, fe = os.path.splitext(f)
    try:
        search[fe].append(f)
    except:
        search[fe]=[f,]

# Example on how to search for your extensions and then do something
extensions = ('.png','.jpg')
for ex in extensions:
    found = search.get(ex,'')
    if found:
        print(found)