2

I want to keep taking input from user, until the user enters a certain string, inside a list comprehension, how to achieve this?

What I need is equivalent to:

lst = []
stop = 'stop'

while True:
    val = input()
    if val == stop:
        break
    lst.append(val)

print('output:',lst)

# Output:
# 1
#
# hello
#
# three
#
# [1, 2, 3]
#
# stop
# output: ['1', 'hello', 'three', '[1, 2, 3]']

But, to reiterate, I don't want while, I would prefer a one-liner, something like:

stop = 'stop'
lst = [i for i in take_input_until(stop)]
print(lst)

# Output:
# 1
#
# hello
#
# three
#
# [1, 2, 3]
#
# stop
# output: ['1', 'hello', 'three', '[1, 2, 3]']

Furthermore, can this be applied to any user defined function, so that it can be calculated until a specific value is returned?

I have checked these questions, which did not meet my requirement:

Sayandip Dutta
  • 15,602
  • 4
  • 23
  • 52

3 Answers3

1

iter to the rescue

Even though iter is mainly used with a single argument, it has a lesser known second argument, that changes how the function behaves.

From python docs of iter:

iter(object[, sentinel])

...If the second argument, sentinel, is given, then object must be a callable object. The iterator created in this case will call object with no arguments for each call to its next() method; if the value returned is equal to sentinel, StopIteration will be raised, otherwise the value will be returned.

One useful application of the second form of iter() is to build a block-reader. For example, reading fixed-width blocks from a binary database file until the end of file is reached:

from functools import partial
with open('mydata.db', 'rb') as f:
   for block in iter(partial(f.read, 64), b''):
       process_block(block)

The above example answers the second question as well. There are other ways to achieve this. In the following examples I would illustrate different ways to take input until a specific condition is met, however, this is just special case, for the second question, it's just that the function is input. So all the examples below, can be made to work for any other function.

This can be used for taking input until a specific input is encountered:

>>> STOP = 'stop'
>>> lst = list(iter(input, STOP))
# can also be written as list comprehension, 
# which would be helpful if you want to do something with the values
#>> lst = [i for i in iter(input, STOP)]

1

hello

three

[1, 2, 3]

stop

>>> print(lst)
['1', 'hello', 'three', '[1, 2, 3]']

Here iter(input, STOP) is something which is known as callable_iterator:

>>> type(iter(input, STOP))
callable_iterator

Showing input prompt

In order to show input prompts for each input, we can use functools.partial:

>>> from functools import partial
>>> lst = [i for i in iter(partial(input, 'enter: '), 'x')]   # or list(iter(partial(input, 'enter: '), 'x'))

enter: 1

enter: 2

enter: 3

enter: x

>>> lst
['1', '2', '3']

Including the stop word

If you wanted to include the stop word in the list as well, you could use iterable unpacking via * operator:

>>> STOP = 'x'
>>> input_with_prompt = partial(input, 'enter: ')
>>> lst = [*iter(input_with_prompt, STOP), STOP]

enter: 1

enter: 2

enter: 3

enter: x

>>> lst
['1', '2', '3', 'x']

This has to be the easiest way to replace while. However, for more complex requirements, this is not very useful.

Sayandip Dutta
  • 15,602
  • 4
  • 23
  • 52
0

takewhile condition is True

We can use a function in itertools.takewhile to check whether the input is equal to stop word, if not, keep taking input:

>>> from itertools import takewhile
>>> STOP = 'x'
>>> lst = list(takewhile(lambda inp: inp != STOP, iter(input_with_prompt, None)))

enter: 1

enter: 2

enter: 3

enter: x

>>> lst
['1', '2', '3']

Here iter(input_with_prompt, None) will keep calling input, because it's sentinel argument will never be met, as input returns str only. This is roughly equivalent to a while True loop, except the values are lazily calculated.

takewhile will call __next__ on the callable_iterator object, and apply the function on the next value, while the condition meets.

This may seem like a redundant and over complicated way of doing the same thing, however, its advantages should be apparent in following examples.

Merits :

  1. takewhile can be used to test for multiple stop_words:
>>> lst = list(takewhile(lambda inp: inp not in ['q', 'quit', 'Q'], iter(input_with_prompt, None)))

enter: 1

enter: 2

enter: 3

enter: quit

>>> lst
['1', '2', '3']
  1. Can be used to pass multiple values to the function, using another iterator type object, here is an example with itertools.count that is used to provide different argument to input at each call.
>>> from itertools import count
>>> lst = list(takewhile(
...     lambda inp: inp != STOP,
...     (input(f'enter val #{i} ("{STOP}" to quit): ') for i in count(1))
... ))

enter val #1 ("x" to quit): 1

enter val #2 ("x" to quit): 2

enter val #3 ("x" to quit): 3

enter val #4 ("x" to quit): x

>>> lst
['1', '2', '3']
  1. Can be combined with range or itertools.repeat to stop either when stop_word is encountered, or number of inputs reaches a specific value.
>>> from itertools import repeat
>>> MAX_NUM = 5
>>> lst = list(takewhile(
...     lambda inp: inp != STOP,
...     ((input('enter : ') for _ in range(MAX_NUM))
... ))

enter : 1

enter : 2

enter : 3

enter : 4

enter : 5
>>> lst
['1', '2', '3', '4', '5']

# ^ Here stop word is not encountered, 
# but input stops when MAX_NUM is reached.
#---------------------------------------------------#
>>> lst = list(takewhile(
...     lambda inp: inp != STOP,
...     (input('enter : ') for _ in repeat(None, MAX_NUM)) 
... ))

enter : 1

enter : 2

enter : 3

enter : x
>>> lst
['1', '2', '3']

# ^ here stop word is encountered before MAX_NUM is reached.

NOTE: (f() for _ in range(n)) behaves the same as (f() for _ repeat(None, n)), however the latter is faster when the when the loop variable is not needed.

  1. Can be combined with itertools.starmap to get nested list for multiple stopwords.
>>> list_of_list = list(list(val) 
...      for val in starmap(
...                     iter, 
...                     [
...                         (input_with_prompt, 'q'), 
...                         (input_with_prompt, 'quit'), 
...                         (input_with_prompt, 'Q')
...                     ])
... ))

enter: 1

enter: 2

enter: q

enter: 2

enter: 3

enter: 4

enter: quit

enter: 5

enter: Q
>>> list_of_list
[['1', '2'], ['2', '3', '4'], ['5']]

While this might look a like a very arcane usecase, it can be especially useful for other domains, for example optimization techniques where you want to check the intermediate steps taken to reach a specific result with different hyperparameters.

So for flexibility, takewhile can be used, but may be at the cost of readability.

Sayandip Dutta
  • 15,602
  • 4
  • 23
  • 52
0

Build your own Iterator

Another option could be making a custom iterator object.

class IterWithCond:
    """Continuously call a function, until a given condition is met."""

    def __init__(
        self, func, cond, include_break=False, max_num=float('inf'), args=(), kwargs={}
    ):
        self.func = func
        self.cond = cond
        self.args = args if isinstance(args, (list, tuple)) else (args,)
        self.kwargs = kwargs
        self.include = include_break
        self.max_num = max_num

    def __iter__(self):
        self._count = 0
        self._cond_met = False
        return self

    def __next__(self):
        if self._cond_met or self._count >= self.max_num:
            raise StopIteration
        else:
            out = self.func(*self.args, **self.kwargs)
            self._cond_met = self.cond(out)
            if not self.include and self._cond_met:
                raise StopIteration
            self._count += 1
            return out
    # Following line enables functionalities like `iter(IterWithCond(*args), stop)`
    __call__ = __next__   

You can choose what configuration you want.

  • Vanila input without prompt, until q or Q is entered.
>>> itr_obj = IterWithCond(
    func=input,
    cond=lambda x: x.lower() == 'q',
)
>>> lst = list(itr_obj)

1

2

3

q
>>> lst
['1', '2', '3']
  • Add prompt, include the break character
>>> itr_obj = IterWithCond(
    func=input,
    cond=lambda x: x.lower() == 'q',
    include_break=True,
    args='enter: '
)
>>> lst = list(itr_obj)

enter: 1

enter: 2

enter: 3

enter: q

>>> lst
['1', '2', '3', 'q']
  • Take until the user inputs 'q'/'Q' or until the user inputs 5 numbers, don't include stop character if it is stopped by stop character.
>>> itr_obj = IterWithCond(
    func=input,
    cond=lambda x: x.lower() == 'q',
    args='enter: ',
    max_num=5
)
>>> lst = list(itr_obj)

enter: 1

enter: 2

enter: 3

enter: 4

enter: 5

>>> lst
['1', '2', '3', '4', '5']

For flexibility and readability, this one is the best. But if it is only worthwhile if it is used multiple times throughout the program with different conditions.

Sayandip Dutta
  • 15,602
  • 4
  • 23
  • 52