How to list all files in a directory with a certain extension in Python

We want to be able to do this:

def main():
    directory = '/home/martin/Python'
    files = list_files(directory, "py")
    for f in files:
        print f
spam.py
eggs.py
ham.py

There are 3 methods we can use:

  1. Use os.listdir()
  2. Use os.walk()
  3. Use glob.glob()

Method1: Use os.listdir()

from os import listdir

def list_files1(directory, extension):
    return (f for f in listdir(directory) if f.endswith('.' + extension))

Method 2: Use os.walk()

Note that walk() will recurse into subdirectories, but we avoid this by returning on the first iteration of the loop.

from os import walk

def list_files2(directory, extension):
    for (dirpath, dirnames, filenames) in walk(directory):
        return (f for f in filenames if f.endswith('.' + extension))

Method 3: Use glob.glob()

You need to change into the directory to use glob(), so it’s good manners to save the current working directory and then change back into it at the end.

from glob import glob
from os import getcwd, chdir

ef list_files3(directory, extension):
    saved = getcwd()
    chdir(directory)
    it = glob('*.' + extension)
    chdir(saved)
    return it