CODEWITHSUNDEEP
pythonstringsortingdirectorylist-comprehension
ProGamerGov
ProGamerGov

Reputation: 970

How do I use list comprehension to simplify the sorting and adding of multiple lists and strings?

The loop below will take a list of directories and/or images, and it will return a list of images, including those images in the supplied directories. For each directory that it is given, it verifies that each file is indeed a valid image, before adding that image to the list.

import os
import argparse
parser = argparse.ArgumentParser()
parser.add_argument("-style_image", help="List of images and/or directories")
params = parser.parse_args()

style_image_input = params.style_image.split(',')
valid_ext, style_input_sorted = [".jpg",".png"], None
for image in style_image_input:
    if os.path.isdir(image):
        for file in os.listdir(image):
            print(file)
            ext = os.path.splitext(file)[1]
            if ext.lower() not in valid_ext:
                continue
            if style_input_sorted == None:
                style_input_sorted = file
            else: 
                style_input_sorted += "," + file
    else:
        if style_input_sorted == None:
            style_input_sorted = image
        else: 
            style_input_sorted += "," + image 
style_image_list = style_input_sorted.split(',')
print(style_image_list)

How can I use list comprehension to simplify this loop?

Upvotes: 0

Views: 75

Answers (1)

abarnert
abarnert

Reputation: 365747

Forget about "how can I turn this into a list comprehension". Start with "how can I simplify this". If, in the end, you get down to a loop with one or two clauses and a simple expression, then you can consider turning that into a list comprehension, but that's the final step, not the starting goal.

Most of your repetition is in the way you're building up style_input_sorted:

  • Set it to None.
  • Every time you get a value, if it's None, set it to the value
  • Otherwise, add a comma and then add the value.

Instead of starting with None, you could start with "", and then do this:

if style_input_sorted:
    style_input_sorted += ","
style_input_sorted += file

But, even more simply: what you're doing is the same thing str.join already knows how to do. If you can build up a list of strings, and then just join that list at the end, it'll be a lot simpler:

style_input_sorted = []
if …
    for …
        style_input_sorted.append(file)
else …
    style_input_sorted.append(file)
style_input_sorted = ",".join(style_input_sorted)

But it looks like the only thing you ever do with style_input_sorted is to split it back out into a list anyway. So why even join a string just to split it?

style_input_list = []
if …
    for …
        style_input_list.append(file)
else …
    style_input_list.append(file)

There are some other simplifications you can make, but this is the biggest one, and it'll open the door for the next ones. For example, now that you're only doing one trivial thing instead of four lines of code for valid extensions, you can probably get rid of the continue:

if os.path.isdir(image):
    for file in os.listdir(image):
        ext = os.path.splitext(file)[1]
        if ext.lower() in valid_ext:
            style_input_list.append(file)
else:
    style_input_list.append(file)

And now we have a piece we could turn into a list comprehension—although I'd use a generator expression:

if os.path.isdir(image):
    images = (file for file in os.listdir(image)
              if os.path.splitext(file)[1].lower() in valid_ext)
    style_input_list.extend(images)
else:
    style_input_list.append(image)

But to turn the whole thing into a list comprehension will be horribly ugly. You can turn that into an expression using ternary if and flatten the whole thing out at the end, but if you have five lines of code, that doesn't belong inside a comprehension. (Of course you could factor those five lines out into a function and then wrap a call to that function into a comprehension, which might be worth doing.)

Upvotes: 3

Related Questions