Handling country codes

1. Find and use well-tested tools

In this example, we'll want to use pycountry, a package that has existed since 2008 and, per its documentation:

…provides the ISO databases for the standards:

• 3166 Countries

The package includes a copy from Debian’s pkg-isocodes and makes the data accessible through a Python API.

It's easy to use and fast:

In [1]: from pycountry import countries

# Get a country by its code
In [2]: countries.get(alpha_2="CA")
Out[2]: Country(alpha_2='CA', alpha_3='CAN', flag='🇨🇦', name='Canada', numeric='124')

# Lookup on any field: alpha_2, alpha_3, or name:
In [3]: countries.lookup("Canada")
Out[3]: Country(alpha_2='CA', alpha_3='CAN', flag='🇨🇦', name='Canada', numeric='124')

# Fuzzy search returns a list() of results
In [4]: countries.search_fuzzy("Korea")
Out[4]: 
[Country(alpha_2='KP', alpha_3='PRK', common_name='North Korea', flag='🇰🇵', name="Korea, Democratic People's Republic of", numeric='408', official_name="Democratic People's Republic of Korea"),
 Country(alpha_2='KR', alpha_3='KOR', common_name='South Korea', flag='🇰🇷', name='Korea, Republic of', numeric='410')]

pycountry is downloaded about 56,000 times every day:

As with pandas, below, heavy usage is both an indicator of and contributor to quality. Hypothetically, we could copy the ISO 3166 table from Wikipedia into a CSV file and write some code to query it. But there's no point in doing this. The result will certainly be slower, more likely to contain bugs, and have fewer features than pycountry.

2. Understand the standard library

Some example data:

In [5]: import pandas as pd

In [6]: data = pd.DataFrame([
   ...:     ["Austria", 2021, 1.1],
   ...:     ["Bosnia", 2021, 2.2],
   ...:     ["Canada", 2021, 3.3],
   ...:     ["UK", 2021, 4.4],
   ...:     ["UK", 2022, 5.5],
   ...: ], columns=["country", "year", "value"])
   ...: 

In [7]: data
Out[7]: 
   country  year  value
0  Austria  2021    1.1
1   Bosnia  2021    2.2
2   Canada  2021    3.3
3       UK  2021    4.4
4       UK  2022    5.5

What is in the "country" column/dimension?

• “Canada” and “Austria” are direct matches for the ‘name’ field of the ISO 3166-1 database.
• “Bosnia” is only a partial match for the name “Bosnia and Herzegovina”.
• “UK” is an idiosyncratic identifier: the ISO 3166-1 codes for the United Kingdom are GB and GBR.

We create a mapping with the idiosyncrasies of this data set:

In [8]: country_map = {
   ...:     "Bosnia": "Bosnia and Herzegovina",
   ...:     "UK": "United Kingdom",
   ...: }
   ...: 

This allows us to look up a correct value for an incorrect one:

In [9]: country_map["UK"]
Out[9]: 'United Kingdom'

In [10]: for name in data["country"]:
   ....:     print(country_map.get(name, name))
   ....: 
Austria
Bosnia and Herzegovina
Canada
United Kingdom
United Kingdom

In [11]: from functools import lru_cache

In [12]: @lru_cache()
   ....: def fix_name(value):
   ....:     return country_map.get(value, value)
   ....: 

In [13]: for name in data["country"]:
   ....:     print(fix_name(name))
   ....: 
Austria
Bosnia and Herzegovina
Canada
United Kingdom
United Kingdom

In [14]: fix_name.cache_info()
Out[14]: CacheInfo(hits=1, misses=4, maxsize=128, currsize=4)

3. Find and use optimized code

Our data is a pandas data structure. Pandas is a very widely used library with many users and contributors. [5] For these reasons, it is internally very sophisticated; many common operations—indeed, almost all operations that most of us will perform in daily use—have been optimized, some using low-level C code, for speed and memory performance.

The simple upshot is that it is almost never necessary to write Python loops (for …: and while …:) when using pandas. If you do find yourself writing loops, it is likely that you can instead find and use existing, optimized functionality in pandas.

I strongly recommend doing this: your code will be both simpler and more performant.

Above, we looped over the values in the "country" column of data. (Remember: each column of a pandas.DataFrame is a pandas.Series.) The feature we should use instead is pandas.Series.apply(). (Again: go read the documentation, all of it.)

This looks like:

In [15]: data["country"].apply(fix_name)
Out[15]: 
0                   Austria
1    Bosnia and Herzegovina
2                    Canada
3            United Kingdom
4            United Kingdom
Name: country, dtype: object

Pandas takes care of calling fix_name() many times, once for each value in this column, and assembling the result into a new series. It internally speeds up and, where possible, parallelizes these calls: the point is we don't need to think about how it does this.

Also, because we've memoized with lru_cache(), after a certain point each call becomes as quick as a dictionary lookup, no matter how expensive fix_name() is:

In [16]: fix_name.cache_info()
Out[16]: CacheInfo(hits=6, misses=4, maxsize=128, currsize=4)

More cache hits have occurred.

Wrapping up

Now we finally bring in pycountry:

In [17]: @lru_cache()
   ....: def code_for_name(value):
   ....:     return countries.get(name=country_map.get(value, value)).alpha_3
   ....: 

This function first uses country_map to correct data set idiosyncrasies, then uses pycountry.get() to look up a record using the ‘name’ field. Finally, it returns the ‘alpha_3’ code from the record.

We apply() this new function to the ‘country’ column:

In [18]: data["country"].apply(code_for_name)
Out[18]: 
0    AUT
1    BIH
2    CAN
3    GBR
4    GBR
Name: country, dtype: object

In [19]: code_for_name.cache_info()
Out[19]: CacheInfo(hits=1, misses=4, maxsize=128, currsize=4)

And we can even update the data, discarding [6] the full names:

In [20]: data.assign(country=data["country"].apply(code_for_name))
Out[20]: 
  country  year  value
0     AUT  2021    1.1
1     BIH  2021    2.2
2     CAN  2021    3.3
3     GBR  2021    4.4
4     GBR  2022    5.5