#!/usr/bin/env python # 3rd party modules from lxml import etree import mpu.pd def main(filename="weekend.svg"): mon_fri = ['Dominican Republic', 'Lebanon', 'Serbia', 'Albania', 'Angola', 'Argentina', 'Armenia', 'Azerbaijan', 'Austria', 'Australia', 'Benin', 'Belarus', 'Belgium', 'Brazil', 'Burundi', 'Bulgaria', 'Canada', 'Cambodia', 'Cameroon', 'Chile', 'China', 'Croatia', 'Costa Rica', 'Czech Republic', 'Denmark', 'Ethiopia', 'Estonia', 'Finland', 'France', 'Gabon', 'Gambia', 'Germany', 'Ghana', 'Greece', 'Hungary', 'Indonesia', 'Ireland', 'Italy', 'Côte d\'Ivoire', 'Japan', 'Kazakhstan', 'Kenya', 'Lao People\'s Democratic Republic', 'Latvia', 'Lesotho', 'Lithuania', 'Madagascar', 'Malawi', 'Mali', 'Malta', 'Mauritania', 'Mongolia', 'Morocco', 'Mozambique', 'Netherlands', 'New Zealand', 'Nigeria', 'Norway', 'Pakistan', 'Poland', 'Portugal', 'Romania', 'Russia', 'Rwanda', 'Senegal', 'Singapore', 'Slovakia', 'Spain', 'Sri Lanka', 'South Africa', 'South Korea', 'Suriname', 'Swaziland', 'Sweden', 'Switzerland', 'Seychelles', 'Taiwan', 'Tanzania', 'Togo', 'Thailand', 'Trinidad and Tobago', 'Tunisia', 'Turkey', 'Ukraine', 'United Kingdom', 'United States', 'Vietnam', 'Congo (Democratic Republic of the)', 'Zambia', 'Zimbabwe'] mon_sat = ['Equatorial Guinea', 'Hong Kong', 'Mexico', 'Philippines', 'Uganda', 'India', 'Colombia'] sun_thu = ['Afghanistan', 'Algeria', 'Bahrain', 'Egypt', 'Iraq', 'Israel', 'Jordan', 'Kuwait', 'Libya', 'Maldives', 'Oman', 'Qatar', 'Saudi Arabia', 'Sudan', 'Syria', 'Yemen', 'United Arab Emirates', 'Bangladesh'] sat_thu = ['Djibouti', 'Palestine', 'Iran'] sun_fri = ['Nepal'] mon_thu_sat = ['Brunei Darussalam'] mixed = ['Malaysia'] countries = mpu.pd.countries countries = countries.set_index('English short name (upper/lower case)') country2dict = countries.to_dict('index') legend = [(mon_fri, '#B8D8D8'), (mon_sat, 'orange'), (sun_thu, '#4F6367'), (sat_thu, 'red'), (sun_fri, 'tomato'), (mon_thu_sat, '#ff00ff'), (mixed, '#00ffff')] colors = {} for collection, color in legend: for country in collection: if country in country2dict: code = country2dict[country]['Alpha-2 code'].lower() colors[code] = color else: print('Could not find the Alpha-2 code for \'{}\' in mpu' .format(country)) tree = read_svg(filename) root = tree.getroot() for country_id in colors.keys(): element = root.xpath('//*[@id="{}"]'.format(country_id)) if len(element) == 1: element = element[0] element.set("style", 'fill: {};'.format(colors[country_id])) for el in element.getchildren(): el.set("style", 'fill: {};'.format(colors[country_id])) else: print("Wrong length for id={}".format(country_id)) write_svg(root) def read_svg(filename): """ Read an SVG file. Parameters ---------- filename : str Returns ------- tree : lxml.ElementTree """ with open(filename, 'r') as f: tree = etree.parse(f) return tree def write_svg(root, filename='weekend-auto.svg'): """ Read an SVG file. Parameters ---------- root : lxml.Element filename : str """ etree.ElementTree(root).write(filename, pretty_print=True) if __name__ == '__main__': main()