NK_
NK_

Reputation: 391

Optimize long if statement in Python

I have a if statement inside a loop that sets a variable to a specific country name if it matches the condition, that is, if the parameter contains the country name.

The parameters are a list of paths that contain a country name in different positions, i.e. C:\\some\\path\\text_country_text.xlsx or C:\\some\\path\\text_text_country_text.xlsx

The if statement is pretty long at the moment because it checks a rather long list of countries. The code I wrote works but it does not look really optimized. Is there a shorter way to do this?

def my_function(*args): 
    for country in args:
        if "Australia" in country:
            country_name = "Australia"
        elif "Austria" in country:
            country_name = "Austria"
        # etc. for many countries

Upvotes: 0

Views: 237

Answers (3)

JohanL
JohanL

Reputation: 6891

Just for the fun of it, and to show the power of sets, it is possible to solve this with only one for loop and a set.intersection() call:

def my_function(*args):
    countries = set(['Austria', 'Australia', 'Brazil', 'Denmark'])
    for country_name in countries.intersection(args):
        print(country_name)

Thus, the loop will only loop over countries that are both among the input arguments and in the predefined countries set.

Upvotes: 1

Susmit
Susmit

Reputation: 394

You could do: Countries = [ all countries] Code:

for country in args:
  for i in Countries:
        if i in country:
          country_name = i

Upvotes: 2

Chris_Rands
Chris_Rands

Reputation: 41168

Assuming you don't know exactly where the country name is within your country string (so you can't use slicing or a regex to extract the relevant portion), you could adapt your approach to take a list of countries, and then use a generator expression with next() to extract the country name:

next((c_name for c_name in countries if c_name in country), 'Unknown')

Examples:

>>> countries = ['Australia', 'Austria']
>>> country = '#Austrian_man#'
>>> next((c_name for c_name in countries if c_name in country), 'Unknown')
'Austria'
>>> country = '#AustrTYPO_man#'
>>> next((c_name for c_name in countries if c_name in country), 'Unknown')
'Unknown'

Upvotes: 4

Related Questions