Script to generate a map of PyTroll contributors

create_contributor_map.py

#!/usr/bin/env python3
# encoding: utf-8
import getpass
import logging
import os
import sys

import json
import geocoder
import numpy as np
from github import Github
from PIL import Image
from pyproj import Proj
from time import sleep
import pycountry
from collections import Counter

from pycoast import ContourWriter, ContourWriterAGG

LOG = logging.getLogger(__name__)
#gcoder = geocoder.arcgis_reverse.ArcgisReverse
#gcoder = geocoder.google
geocoders = [geocoder.google, geocoder.arcgis_reverse.ArcgisReverse]
MAX_CONTRIBUTORS = 10.


def get_all_pytroll_contributors(user):
    g = Github(user, getpass.getpass('GitHub Password: '))
    pytroll_org = g.get_organization('pytroll')
    LOG.info("Getting all PyTroll repositories...")
    pytroll_repos = [x for x in pytroll_org.get_repos() if x.name != u'aggdraw']
    LOG.info("Getting all PyTroll contributors...")
    all_pytroll_contributors = [
        u for r in pytroll_repos for u in r.get_contributors()]
    set_pytroll_contributors = {u.login: u for u in all_pytroll_contributors}
    return set_pytroll_contributors


def get_country(locations, cache_file='countries.json'):
    """Return 3-digit country code for each location provided.
    """
    # cache country results because most APIs have query limits
    if os.path.isfile(cache_file):
        json_cache = json.load(open(cache_file, 'r'))
    else:
        json_cache = {}

    need_sleep = False
    for loc in locations:
        for gcoder in geocoders:
            if loc in json_cache:
                yield json_cache[loc]
                need_sleep = False
                break
            else:
                need_sleep = True
                LOG.debug("Finding: %s", loc)
                result = gcoder(loc)
                if not result.ok or result.country is None:
                    LOG.debug("Bad result using %r, country %r", gcoder, result.country)
                    continue
                country = result.country
                if len(country) == 2:
                    country = pycountry.countries.get(alpha_2=country).alpha_3
                else:
                    try:
                        country = pycountry.countries.get(alpha_3=country).alpha_3
                    except KeyError:
                        LOG.error("Invalid country code: %s", country)
                        continue
                if country is not None:
                    json_cache[loc] = country
                yield country
                break
        else:
            try:
                # last resort
                country = pycountry.countries.get(name=loc).alpha_3
                json_cache[loc] = country
                yield country
                continue
            except KeyError:
                pass
            LOG.warning("Could not find country for {}".format(loc))
            yield None
        if need_sleep:
            sleep(1.0)

    json.dump(json_cache, open(cache_file, 'w'), indent=4, sort_keys=True)


def main():
    import argparse
    parser = argparse.ArgumentParser(
        description="Create a world map image of PyTroll contributor locations")
    parser.add_argument('--map-proj', default="+proj=moll",
                        help='PROJ.4 map projection string to use for the output image')
    parser.add_argument('--output-size', default=(800, 600), nargs=2, type=int,
                        help='\'width height\' of output image')
    parser.add_argument('--contributor-color', default=(255, 127, 127),
                        help='Color of dots for each contributor')
    parser.add_argument('--bg-color', default=None, nargs=3, type=int,
                        help='Background color of image \'R G B\' 0-255 (default: transparent)')
    parser.add_argument('-o', '--output', default='contributors_map.png',
                        help='Output filename to save the image to')
    parser.add_argument('-u', '--github-user', required=True,
                        help='github username')
    parser.add_argument('gshhg_dir',
                        help='root directory for GHSSG shape files')
    args = parser.parse_args()

    logging.basicConfig(level=logging.INFO)

    im = Image.new('RGBA', args.output_size, color=tuple(args.bg_color) if args.bg_color is not None else None)
    cw = ContourWriterAGG(args.gshhg_dir)
    p = Proj(args.map_proj)
    a, _ = p(-180, 0)
    _, b = p(0, -90)
    c, _ = p(180, 0)
    _, d = p(0, 90)

    extents = (a, b, c, d)
    area_def = (args.map_proj, extents)

    LOG.info("Getting PyTroll contributors...")
    contributors = get_all_pytroll_contributors(args.github_user)
    LOG.info("Getting all PyTroll contributor locations...")
    all_user_locations = []
    for user in contributors.values():
        if user.location is None or user.location == '':
            LOG.info("User location not specified:  ID: '{}';\tName: '{}';\tURL: '{}'".format(user.id, user.name, user.url))
            continue
        # country = next(get_country([user.location]))
        # if country is None:
        #     LOG.info("Unknown contributor location: %s", user.location)
        #     continue
        all_user_locations.append(user.location)
        # number_contributors[country] += 1
        # print(user.login, user.location, country, number_contributors[country])
    all_countries = list(get_country(all_user_locations))
    number_contributors = Counter(all_countries)
    del number_contributors[None]
    all_countries = list(number_contributors.keys())

    def shape_generator():
        import shapefile
        reader = shapefile.Reader("ne_110m_admin_0_countries.shp")
        name_idx = [idx for idx, f in enumerate(reader.fields) if f[0] == 'name'][0]
        new_kwargs = dict(
                fill=args.contributor_color,
                outline=(0, 0, 0),
                fill_opacity=200,
                outline_opacity=255,
                width=1)
        default_kwargs = dict(
                fill=(127, 127, 127),
                outline=(0, 0, 0),
                fill_opacity=127,
                outline_opacity=255,
                width=1)

        class FakeShape(object):
            def __init__(self, points, bbox):
                self.points = points
                self.bbox = bbox

        for sr in reader.shapeRecords():
            name = sr.record[name_idx]
            if not isinstance(name, str):
                try:
                    name = name.decode()
                except UnicodeDecodeError:
                    name = name.decode('latin-1')
            shape_country = list(get_country([name]))[0]
            shape_parts = []
            for part_idx, part_idx2 in zip(sr.shape.parts, list(sr.shape.parts[1:]) + [len(sr.shape.points)]):
                shape_parts.append(FakeShape(sr.shape.points[part_idx: part_idx2], sr.shape.bbox))

            if shape_country is not None and shape_country in all_countries:
                LOG.info("Contributor country: %s, Code: %s, Number of Contributors: %d", name, shape_country, number_contributors[shape_country])
                kwargs = new_kwargs.copy()
                num_cont = min(number_contributors[shape_country], MAX_CONTRIBUTORS)
                kwargs['fill_opacity'] = int((num_cont / 10.) * 55 + 200)
                yield shape_parts, kwargs
            else:
                LOG.debug("No contributor's from Country: %s, Code: %s", name, shape_country)
                yield shape_parts, default_kwargs

    LOG.info("Applying contributor locations to map image...")
    cw._add_feature(im, area_def, 'polygon', 'custom', shape_generator=shape_generator())
    im.save(args.output, dpi=(300, 300))


if __name__ == "__main__":
    sys.exit(main())

pip install pygithub
pip install geocoder # or conda install geocoder