'NoneType' object has no attribute 'name' when using the pycountry library

Question

I have a column in a DataFrame with country code and I want to convert it to the name to plot a graph using the library pycountry

def get_country(n):
    country = countries.get(alpha_2 = n)
    return country.name

I want to implement the function above using on the DataFrame like this

df['country'] = df['country'].apply(get_country)

and I get this error

AttributeError: 'NoneType' object has no attribute 'name'

Answer 1

get() returns None by default if the key isn't found. You need to check for this.

def get_country(n):
    country = countries.get(alpha_2 = n)
    if country:
        return country.name
    else:
        return n # keep the original code if no name found