PParker
PParker

Reputation: 1511

How to avoid very long if-elif-elif-else statements in Python function

Is there a smart way to shorten very long if-elif-elif-elif... statements?

Let's say I have a function like this:

def very_long_func():
  something = 'Audi'
  
  car = ['VW', 'Audi', 'BMW']
  drinks = ['Cola', 'Fanta', 'Pepsi']
  countries = ['France', 'Germany', 'Italy']
  
  if something in car:
    return {'type':'car brand'}
  elif something in drinks:
    return  {'type':'lemonade brand'}
  elif something in countries:
    return {'type':'country'}
  else:
    return {'type':'nothing found'}
  

very_long_func()
>>>> {'type': 'car brand'}

The actual function is much longer than the example. What would be the best way to write this function (not in terms of speed but in readability)

I was reading this, but I have trouble to apply it to my problem.

Upvotes: 0

Views: 1070

Answers (4)

Alain T.
Alain T.

Reputation: 42133

You can simulate a switch statement with a helper function like this:

def switch(v): yield lambda *c: v in c

The your code could be written like this:

something = 'Audi'

for case in switch(something):
    if case('VW', 'Audi', 'BMW'):          name = 'car brand'      ; break
    if case('Cola', 'Fanta', 'Pepsi'):     name = 'lemonade brand' ; break
    if case('France', 'Germany', 'Italy'): name = 'country'        ; break
else:                                      name = 'nothing found'
return {'type':name}

If you don't have specific code to do for each value, then a simple mapping dictionary would probably suffice. For ease of maintenance, you can start with a category-list:type-name mapping and expand it before use:

mapping = { ('VW', 'Audi', 'BMW'):'car brand',
            ('Cola', 'Fanta', 'Pepsi'):'lemonade brand',
            ('France', 'Germany', 'Italy'):'country' }
mapping = { categ:name for categs,name in mapping.items() for categ in categs }

Then your code will look like this:

something = 'Audi'
return {'type':mapping.get(something,'nothing found')}

using a defaultdict would make this even simpler to use by providing the 'nothing found' value automatically so you could write: return {'type':mapping[something]}

Upvotes: 0

Pygirl
Pygirl

Reputation: 13349

you can create dictionary like this and then use map_dict.

from functools import reduce
car = ['VW', 'Audi', 'BMW']
drinks = ['Cola', 'Fanta', 'Pepsi']
countries = ['France', 'Germany', 'Italy']

li = [car, drinks, countries]
types = ['car brand', 'lemonade brand', 'country', 'nothing found']
dl = [dict(zip(l, [types[idx]]*len(l))) for idx, l in enumerate(li)]

map_dict = reduce(lambda a, b: dict(a, **b), dl)

Upvotes: 2

dimay
dimay

Reputation: 2804

Try this:

def create_dct(lst, flag):
    return {k:flag for k in lst}

car = ['VW', 'Audi', 'BMW']
drinks = ['Cola', 'Fanta', 'Pepsi']
countries = ['France', 'Germany', 'Italy']

merge_dcts = {}
merge_dcts.update(create_dct(car, 'car brand'))
merge_dcts.update(create_dct(drinks, 'lemonade brand'))
merge_dcts.update(create_dct(countries, 'country'))

something = 'Audi'
try:
    print("type: ", merge_dcts[something])
except:
    print("type: nothing found")

Upvotes: 1

AnkurSaxena
AnkurSaxena

Reputation: 825

You can't hash lists as dictionary values. So go other way round. Create a mapping of type -> list. And initialize your output with the default type. This allows you to keep on adding new types to your mapping without changing any code.

def very_long_func():
  something = 'Audi'
  
  car = ['VW', 'Audi', 'BMW']
  drinks = ['Cola', 'Fanta', 'Pepsi']
  countries = ['France', 'Germany', 'Italy']
  
  out = {'type': 'nothing found'}  # If nothing matches
  mapping = {
      'car brand': car,
      'lemonade brand': drinks,
      'country': countries
    }
  for k,v in mapping.items() :
    if something in v:
      out['type'] = k    # update if match found
      break
  return out             # returns matched or default value

Upvotes: 3

Related Questions