User:TolBot/Task 5

run = True

languages = {
	'en': {
		'page': 'Template:COVID-19 data/data',
		'summary': '[[m:User:TolBot/Task 5|Task 5]]: update COVID-19 pandemic data'
	},
	'fr': {
		'page': 'Modèle:Données de la pandémie de Covid-19/données',
		'summary': '[[m:User:TolBot/Task 5|Tâche 5]]: mettre à jour les données de la pandémie de COVID-19'
	},
	'et': {
		'page': 'Mall:Koroonapandeemia riigiti/andmed',
		'summary': '[[m:User:TolBot/Task 5|Ülesanne 5]]: värskendada COVID-19 pandeemia andmeid'
	},
	'af': {
		'page': 'Sjabloon:Koronaviruspandemie van 2019-2020/data',
		'summary': '[[m:User:TolBot/Task 5|Taak 5]]: dateer COVID-19-pandemiedata op'
	},
	'uk': {
		'page': 'Шаблон:Пандемія COVID-19 за країнами та територіями/дані',
		'summary': '[[m:User:TolBot/Task 5|Завдання 5]]: оновити дані про пандемію COVID-19'
	},
	'zh': {
		'page': 'Template:COVID-19 data/data',
		'summary': '[[m:User:TolBot/Task 5|任務5]]：更新 COVID-19 疫情数据'
	}
}

bot_username = ''
bot_password = ''
bot_credentials = (bot_username, bot_password)

import requests
import json
import pandas as pd
import re

api = 'https://{0}.wikipedia.org/w/api.php?format=json&formatversion=2'

owid_url = 'https://catalog.ourworldindata.org/garden/covid/latest/compact/compact.csv'

owid_translator = {
	'countries': {
		'OWID_WRL': 'XW',
		'OWID_EU27': 'EU',
		'OWID_CYN': 'XC', # Northern Cyprus
		'OWID_KOS': 'XK', # Kosovo
		'OWID_TRS': 'XT', # Transnistria
		'ABW': 'AW',
		'AFG': 'AF',
		'AGO': 'AO',
		'AIA': 'AI',
		'ALA': 'AX',
		'ALB': 'AL',
		'AND': 'AD',
		'ARE': 'AE',
		'ARG': 'AR',
		'ARM': 'AM',
		'ASM': 'AS',
		'ATA': 'AQ',
		'ATF': 'TF',
		'ATG': 'AG',
		'AUS': 'AU',
		'AUT': 'AT',
		'AZE': 'AZ',
		'BDI': 'BI',
		'BEL': 'BE',
		'BEN': 'BJ',
		'BES': 'BQ',
		'BFA': 'BF',
		'BGD': 'BD',
		'BGR': 'BG',
		'BHR': 'BH',
		'BHS': 'BS',
		'BIH': 'BA',
		'BLM': 'BL',
		'BLR': 'BY',
		'BLZ': 'BZ',
		'BMU': 'BM',
		'BOL': 'BO',
		'BRA': 'BR',
		'BRB': 'BB',
		'BRN': 'BN',
		'BTN': 'BT',
		'BVT': 'BV',
		'BWA': 'BW',
		'CAF': 'CF',
		'CAN': 'CA',
		'CCK': 'CC',
		'CHE': 'CH',
		'CHL': 'CL',
		'CHN': 'CN',
		'CIV': 'CI',
		'CMR': 'CM',
		'COD': 'CD',
		'COG': 'CG',
		'COK': 'CK',
		'COL': 'CO',
		'COM': 'KM',
		'CPV': 'CV',
		'CRI': 'CR',
		'CUB': 'CU',
		'CUW': 'CW',
		'CXR': 'CX',
		'CYM': 'KY',
		'CYP': 'CY',
		'CZE': 'CZ',
		'DEU': 'DE',
		'DJI': 'DJ',
		'DMA': 'DM',
		'DNK': 'DK',
		'DOM': 'DO',
		'DZA': 'DZ',
		'ECU': 'EC',
		'EGY': 'EG',
		'ERI': 'ER',
		'ESH': 'EH',
		'ESP': 'ES',
		'EST': 'EE',
		'ETH': 'ET',
		'FIN': 'FI',
		'FJI': 'FJ',
		'FLK': 'FK',
		'FRA': 'FR',
		'FRO': 'FO',
		'FSM': 'FM',
		'GAB': 'GA',
		'GBR': 'GB',
		'GEO': 'GE',
		'GGY': 'GG',
		'GHA': 'GH',
		'GIB': 'GI',
		'GIN': 'GN',
		'GLP': 'GP',
		'GMB': 'GM',
		'GNB': 'GW',
		'GNQ': 'GQ',
		'GRC': 'GR',
		'GRD': 'GD',
		'GRL': 'GL',
		'GTM': 'GT',
		'GUF': 'GF',
		'GUM': 'GU',
		'GUY': 'GY',
		'HKG': 'HK',
		'HMD': 'HM',
		'HND': 'HN',
		'HRV': 'HR',
		'HTI': 'HT',
		'HUN': 'HU',
		'IDN': 'ID',
		'IMN': 'IM',
		'IND': 'IN',
		'IOT': 'IO',
		'IRL': 'IE',
		'IRN': 'IR',
		'IRQ': 'IQ',
		'ISL': 'IS',
		'ISR': 'IL',
		'ITA': 'IT',
		'JAM': 'JM',
		'JEY': 'JE',
		'JOR': 'JO',
		'JPN': 'JP',
		'KAZ': 'KZ',
		'KEN': 'KE',
		'KGZ': 'KG',
		'KHM': 'KH',
		'KIR': 'KI',
		'KNA': 'KN',
		'KOR': 'KR',
		'KWT': 'KW',
		'LAO': 'LA',
		'LBN': 'LB',
		'LBR': 'LR',
		'LBY': 'LY',
		'LCA': 'LC',
		'LIE': 'LI',
		'LKA': 'LK',
		'LSO': 'LS',
		'LTU': 'LT',
		'LUX': 'LU',
		'LVA': 'LV',
		'MAC': 'MO',
		'MAF': 'MF',
		'MAR': 'MA',
		'MCO': 'MC',
		'MDA': 'MD',
		'MDG': 'MG',
		'MDV': 'MV',
		'MEX': 'MX',
		'MHL': 'MH',
		'MKD': 'MK',
		'MLI': 'ML',
		'MLT': 'MT',
		'MMR': 'MM',
		'MNE': 'ME',
		'MNG': 'MN',
		'MNP': 'MP',
		'MOZ': 'MZ',
		'MRT': 'MR',
		'MSR': 'MS',
		'MTQ': 'MQ',
		'MUS': 'MU',
		'MWI': 'MW',
		'MYS': 'MY',
		'MYT': 'YT',
		'NAM': 'NA',
		'NCL': 'NC',
		'NER': 'NE',
		'NFK': 'NF',
		'NGA': 'NG',
		'NIC': 'NI',
		'NIU': 'NU',
		'NLD': 'NL',
		'NOR': 'NO',
		'NPL': 'NP',
		'NRU': 'NR',
		'NZL': 'NZ',
		'OMN': 'OM',
		'PAK': 'PK',
		'PAN': 'PA',
		'PCN': 'PN',
		'PER': 'PE',
		'PHL': 'PH',
		'PLW': 'PW',
		'PNG': 'PG',
		'POL': 'PL',
		'PRI': 'PR',
		'PRK': 'KP',
		'PRT': 'PT',
		'PRY': 'PY',
		'PSE': 'PS',
		'PYF': 'PF',
		'QAT': 'QA',
		'REU': 'RE',
		'ROU': 'RO',
		'RUS': 'RU',
		'RWA': 'RW',
		'SAU': 'SA',
		'SDN': 'SD',
		'SEN': 'SN',
		'SGP': 'SG',
		'SGS': 'GS',
		'SHN': 'SH',
		'SJM': 'SJ',
		'SLB': 'SB',
		'SLE': 'SL',
		'SLV': 'SV',
		'SMR': 'SM',
		'SOM': 'SO',
		'SPM': 'PM',
		'SRB': 'RS',
		'SSD': 'SS',
		'STP': 'ST',
		'SUR': 'SR',
		'SVK': 'SK',
		'SVN': 'SI',
		'SWE': 'SE',
		'SWZ': 'SZ',
		'SXM': 'SX',
		'SYC': 'SC',
		'SYR': 'SY',
		'TCA': 'TC',
		'TCD': 'TD',
		'TGO': 'TG',
		'THA': 'TH',
		'TJK': 'TJ',
		'TKL': 'TK',
		'TKM': 'TM',
		'TLS': 'TL',
		'TON': 'TO',
		'TTO': 'TT',
		'TUN': 'TN',
		'TUR': 'TR',
		'TUV': 'TV',
		'TWN': 'TW',
		'TZA': 'TZ',
		'UGA': 'UG',
		'UKR': 'UA',
		'UMI': 'UM',
		'URY': 'UY',
		'USA': 'US',
		'UZB': 'UZ',
		'VAT': 'VA',
		'VCT': 'VC',
		'VEN': 'VE',
		'VGB': 'VG',
		'VIR': 'VI',
		'VNM': 'VN',
		'VUT': 'VU',
		'WLF': 'WF',
		'WSM': 'WS',
		'YEM': 'YE',
		'ZAF': 'ZA',
		'ZMB': 'ZM',
		'ZWE': 'ZW'
	},
	'columns': {
		'people_vaccinated': 'total_vaccinated',
		'people_fully_vaccinated': 'fully_vaccinated',
		'total_vaccinations': 'vaccine_doses',
		'total_cases': 'cases',
		'total_deaths': 'deaths',
		'people_vaccinated_per_hundred': 'percent_vaccinated',
		'people_fully_vaccinated_per_hundred': 'percent_fully_vaccinated',
		'total_deaths_per_million': 'deaths_per_million',
		'population': 'population'
	}
}

# Login
def login(api, session, credentials):
  token_params = {
      'action': 'query',
      'meta': 'tokens',
      'type': 'login'
  }
  token_q = session.post(api, data=token_params).json()
  token = token_q['query']['tokens']['logintoken']
  username, password = credentials
  login_params = {
      'action': 'login',
      'lgname': username,
      'lgpassword': password,
      'lgtoken': token
  }
  login_q = session.post(api, data=login_params).json()
  return session

# Get and parse current data on page
def parse_json(api, session, page):
  params = {
      'action': 'parse',
      'prop': 'wikitext',
      'page': page
  }
  q = session.post(api, data=params).json()
  wikitext = q['parse']['wikitext']
  data = json.loads(wikitext)
  return data

def get(url):
  return pd.read_csv(url).groupby('code').max('date')

# Update with OWID data
def update_owid(owid_data, data, translator):
  for owid_code, location_data in owid_data.iterrows():
    if owid_code in translator['countries']:
      code = translator['countries'][owid_code]
      for owid_col, col in translator['columns'].items():
        if owid_col in location_data:
          if not code in data:
            data[code] = {}
          if not pd.isna(location_data[owid_col]):
            data[code][col] = location_data[owid_col]
  return data

# Old update_owid() for reference, used with:
# https://covid.ourworldindata.org/data/latest/owid-covid-latest.json
# (broken now that GitHub repo not updated)
#
# def update_owid(owid_data, data, translator):
#   for owid_code, location_data in owid_data.items():
#     if owid_code in translator['countries']:
#       for owid_col, col in translator['columns'].items():
#         if owid_col in location_data:
#           code = translator['countries'][owid_code]
#           if not code in data:
#             data[code] = {}
#           if location_data[owid_col]:
#             data[code][col] = location_data[owid_col]
#   return data

# Edit page with updated data
def edit(api, session, page, data, summary):
  text = json.dumps(data, ensure_ascii=False, sort_keys=True)
  token_params = {
      'action': 'query',
      'meta': 'tokens'
  }
  token_q = session.post(api, data=token_params).json()
  token = token_q['query']['tokens']['csrftoken']
  edit_params = {
      'action': 'edit',
      'title': page,
      'text': text,
      'summary': summary,
      'token': token
  }
  edit_q = session.post(api, data=edit_params).json()
  return edit_q

session = requests.Session()
if run:
  for lang in languages:
    login(api.format(lang), session, bot_credentials)
owid_data = get(owid_url)
for lang, lang_data in languages.items():
  data = parse_json(api.format(lang), session, lang_data['page'])
  data = update_owid(owid_data, data, owid_translator)
  if run:
    print(lang, edit(
      api.format(lang),
      session,
      lang_data['page'],
      data,
      lang_data['summary']
    ))
  else:
    print(
      '\n\n', lang, '\n',
      json.dumps(data, indent=2, ensure_ascii=False, sort_keys=True)
    )