Prepend the 3-letter ISO country code and continent name to every line, except the 1st line, of a csv file containing country names

country-codes.csv

country-encode.py

#!/usr/bin/python3

# Prepend the 3-letter ISO country code and continent name to every line,
# except the 1st line, of a owid csv file containing country names

# 幫 OWID 資料加上三字元 iso 國碼及所在大陸
# https://newtoypia.blogspot.com/2023/04/owid-iso.html

# This is a revised version after the discussion in
# https://github.com/owid/owid-datasets/issues/4

# This program requires the presence of country-codes.csv
# in the same directory, or specified as a command line option.
# This reference csv file is a "join" of the following two tables.
# Country or region names follow owid conventions.
# https://statisticstimes.com/geography/countries-by-continents.php
# https://catalog.ourworldindata.org/garden/regions/2023-01-01/regions/definitions.feather

# Sample data file democracy.csv is taken from
# https://github.com/owid/owid-datasets/tree/master/datasets/Democracy%20-%20Economist%20Intelligence%20Unit%20(2023)
# You can also try this .csv for some failed matches:
# https://github.com/owid/owid-datasets/tree/master/datasets/Electricity%20mix%20from%20BP%20%26%20EMBER%20(2022)

# Output may contain lines beginning with #, indicating
# that the program fails to find a unique mapping.
# In such cases further manual adjustment on the input csv file
# is required before feeding it to this program.

# usage: python3 country-encode.py democracy.csv

import argparse, csv, re, difflib

parser = argparse.ArgumentParser(
    description='replace country names by short codes',
    formatter_class=argparse.ArgumentDefaultsHelpFormatter)
parser.add_argument('-t', '--type', type=str, default='iso3',
    help='type of short codes, either iso3 or m49')
parser.add_argument('-c', '--col', type=int, default=0,
    help='column of country name in input csv file')
parser.add_argument('--reference', type=str, default='country-codes.csv',
    help='path to reference file')
parser.add_argument('csvfile', help='csv file')
args = parser.parse_args()

def shorten(name):
    name = re.sub('(people''s )?republic( of)?', '', name, flags=re.I)
    name = re.sub('region', '', name, flags=re.I)
    return name.strip()

by_owid_name = {}
by_other_name = {}
by_short_name = {}
with open(args.reference) as F:
    reftab = csv.DictReader(F)
    for row in reftab:
        row = { k:row[k] for k in row }
        # create a new copy
        by_owid_name[row['owid_name']] = row
        if 'other_name' in row: by_other_name[row['other_name']] = row
        by_short_name[shorten(row['owid_name'])] = row

# print(len(by_other_name), len(by_owid_name), len(by_short_name))
empty = { k:'' for k in by_owid_name[list(by_owid_name.keys())[0]] }
empty['iso3'] = '#' ; empty['m49'] = '#'

short_names = by_short_name.keys()
# for n in short_names: print('<{}>'.format(n))

with open(args.csvfile) as F:
    data = csv.reader(F)
    header = next(data)
    print(','.join([args.type, 'continent'] + header))
    for fields in data:
        cname = fields[args.col]
        if cname in by_owid_name:
            match = by_owid_name[cname]
        elif cname in by_other_name:
            match = by_other_name[cname]
        elif shorten(cname) in by_short_name:
            match = by_short_name[shorten(cname)]
        else:
            ccode = difflib.get_close_matches(shorten(cname), short_names)
            if len(ccode) == 1:
                match = by_short_name[ccode[0]]
            elif len(ccode) > 1:
                match = { k:match[k] for k in match }
                match[args.type] = '# ' + shorten(cname) + ': ' + '|'.join( [ by_short_name[cc][args.type] for cc in ccode ] )
            else:
                match = empty
        print( ','.join( [match[args.type], match['continent']] + fields ) )

democracy.csv