How to check the function condition inside the list comprehension in Python

Question

Hi i was trying to find out if there is a way to do it.

number_list = [d(x) for x in range(1, 10001) if d(x) < 10000]

I want to use list comprehension something like this.

So basically, I want to have a list items with the function result that is less than 10000.

Is there a way to check the condition of the function and put it into the list by using list comprehension?

Answer 1

start with a generator

a_generator = (d(x) for x in a_list)
# a generator will only evaluate when iterated
my_list = [v for v in a_generator if v < 1000]
# all values should be less than 1000 in this example
max(my_list) # better be less than 1000

if the values in a_list are sorted from low to high and the rule of

d(x) < d(x + 1) is followed you can optimize it further

import itertools
a_generator = (d(x) for x in my_sorted_list)
my_list = list(itertools.takewhile(lambda v:v < 1000, a_generator))
# this one will only calculate while the result meets the condition
# once a single value exceeds it will stop taking values

you can of coarse easily transform either of these into a single line

[v for v in (d(x) for x in a_list) if v < 1000]

and

list(itertools.takewhile(lambda v:v < 1000,  (d(x) for x in my_sorted_list)))

respectively, but future you and anyone you work with would probably prefer it on 2 lines

Answer 2

You can write comprehension as follows:

res = [d(x) if d(x) < 10000 else None for x in range(1, 10001)]
number_list = [i for i in res if i is not None]

List comprehensions can have two spots where if/else can be used:

foo = [{x if <smth> else y if <smth_else> else z} for i in range(1000) {if i % 2 == 0}]

The first if-else (inside first pair of curly brackets) determines what exactly to put in your list and the second if statement (second pair of curly brackets) determines when to do it.

Note: curly brackets are for display purposes only. They shouldn't be used in the actual code.

For example a list comprehension of FizzBuzz would look like this:

fb = ['FizzBuzz' if i % 15 == 0 
      else 'Fizz' if i % 3 == 0 
      else 'Buzz' if i % 5 == 0 
      else str(i) 
      for i in range(1, 100)]

EDIT: Updated as per Joran's comment. I forgot that else statement is necessary when using if-else in the front part of comprehension.