Search code examples
pythonlistlist-comprehensionglob

List comprehension with if conditional to get list of files of a specific type


Roughly following the logic for putting together list comprehensions with an if-else described in this solution, I am trying to get a list of files with a specific extensions under a parent directory.

Here is the long form of the code:

mxd_list = []
for top_dir, dir_list, obj_list in os.walk(top_path):
    for obj in obj_list:
        if obj.endswith('.mxd'):
            mxd_list.append(os.path.join(top_dir, obj))

This is my current attempt at consolidating this using a list comprehension. Although it runs, the list is empty.

for top_dir, dir_list, obj_list in os.walk(top_path):
    mxd_list = [os.path.join(top_dir, obj) for obj in obj_list if obj.endswith('.mxd')]

Solution

  • You were quite near. You need to append to the list which is outside the loop

    mxd_list = []
    for top_dir, dir_list, obj_list in os.walk(top_path):
        mxd_list.extend([os.path.join(top_dir, obj) for obj in obj_list if obj.endswith('.mxd')])
    

    The error is that - In every iteration of the outer for loop, the list comp would generate a list specific to that iteration only, thus you need to extend each of the list generated after every iteration to the outer variable, mxd_list.

    Note - The [ are redundant in the sense that removing them would make the contents a generator expression. That is the the statement can be written as mxd_list.extend(os.path.join(top_dir, obj) for obj in obj_list if obj.endswith('.mxd')).

    Another Way of doing it - Using glob.iglob

    As Padraic mentions

    mxd_list = []
    for top_dir, dir_list, obj_list in os.walk(top_path):
         mxd_list.extend(iglob(top_dir+"/*.mxd"))
    

    Is a better way of doing it. But do not forget to import the module, i.e, from glob import iglob