User:Yinweichen/exo.py

# -*- coding: utf-8 -*-
# Winston Yin 2015-07-25
# Takes in (currently only) NASA Exoplanet Archive csv file
# and parses it into a sortable wikitable.
# Select only 12 columns to export.

import pywikibot
import csv
import re
from math import log10, floor

# Basically one row of the table
class exoplanet:
	def __init__(self, row):
		self.n = row[0] # Always use as integer
		self.host = row[1]
		self.letter = row[2]
		self.method = row[3]
		self.period = tofloat(row[5])
		self.a = tofloat(row[6])
		self.ecc = tofloat(row[7])
		self.mass = tofloat(row[8])
		self.radius = tofloat(row[9])
		self.incl = tofloat(row[10])
		self.dist = tofloat(row[11])
		self.year = row[12]
		self.numpl = int(row[4])

	# Wikitext of one row
	def __str__(self):
		row = [None]*12
		row[0] = link(self.host, self.letter)
		row[1] = '[[' + translate(self.host) + ']]'
		row[2] = sigfig(self.radius, 4)
		row[3] = sigfig(self.mass, 4)
		row[4] = sigfig(self.period, 4)
		row[5] = sigfig(self.a, 4)
		row[6] = sigfig(self.ecc, 4)
		row[7] = sigfig(self.incl, 4)
		row[8] = sigfig(self.dist, 4)
		row[9] = '[[' + translate(self.method) + ']]'
		row[10] = self.year
		row[11] = self.numpl
		string = '|-'
		if row[11] > 1:
			string += ' style=\"background:#faecc8;\"'
		string += '\n| ' + ' || '.join(row[0:11]) + '\n'
		return string

# Load csv file into a list of exoplanets
def load_csv(filename):
	table = []
	with open(filename, 'r') as csvfile:
		csvrows = csv.reader(csvfile)
		for row in csvrows:
			exo = exoplanet(row)
			table.append(exo)
	print(str(len(table)) + ' exoplanets in total.')
	return table

# Generates Wikitext for entire table
def wikify(table):
	outtext = ('{| class=\"wikitable sortable\"\n'
		+ '! [[行星]] !! 主星 !! 半徑<br /><small>(×[[木星]]半徑)</small> !! '
		+ '質量<br /><small>(×[[木星]]質量)</small> !! [[軌道周期]]<br /><small>(天)</small> !! '
		+ '軌道[[半長軸]]<br /><small>([[天文單位|AU]])</small> !! [[軌道離心率]] !! [[軌道傾角]] !! '
		+ '距離<br /><small>([[秒差距|pc]])</small> !! [[系外行星偵測法|發現方法]] !! 發現年份\n')
	for exo in table:
		outtext += exo.__str__()
	outtext += '|}'
	return outtext

# Translates links
def translate(string):
	lib = [
		['Radial Velocity', '徑向速度法'],
		['Imaging', '直接拍照法'],
		['Eclipse Timing Variations', '凌日時間變分法'],
		['Transit Timing Variations', '凌日時間變分法'], # Different?
		['Transit', '凌日法'],
		['Astrometry', '天體測量法'],
		['Microlensing', '重力微透鏡法'],
		['Orbital Brightness Modulation', '軌道相位光度變化法'],
		['Pulsar Timing', '脈衝星計時法'],
		['Pulsation Timing Variations', '脈動時間變分法'], # Different?
		[r'Kepler-(\d+) *(|[A-Z])', r'開普勒\1\2'],
		['Lupus', '豺狼座'],
		['gam 1 Leo', '獅子座γ'],
		[r'(\w+) Lup *(|[A-Z])', r'豺狼座\1\2'],
		[r'(\w+) Com *(|[A-Z])', r'后髮座\1\2'],
		[r'(\w+) UMi *(|[A-Z])', r'小熊座\1\2'],
		[r'(\w+) And *(|[A-Z])', r'仙女座\1\2'],
		[r'(\w+) Her *(|[A-Z])', r'武仙座\1\2'],
		[r'(\w+) UMa *(|[A-Z])', r'大熊座\1\2'],
		[r'(\w+) Lyn *(|[A-Z])', r'天貓座\1\2'],
		[r'(\w+) CMa *(|[A-Z])', r'大犬座\1\2'],
		[r'(\w+) Cyg *(|[A-Z])', r'天鵝座\1\2'],
		[r'(\w+) Del *(|[A-Z])', r'海豚座\1\2'],
		[r'(\w+) Sex *(|[A-Z])', r'六分儀座\1\2'],
		[r'(\w+) Ari *(|[A-Z])', r'白羊座\1\2'],
		[r'(\w+) Dra *(|[A-Z])', r'天龍座\1\2'],
		[r'(\w+) Peg *(|[A-Z])', r'飛馬座\1\2'],
		[r'(\w+) Cnc *(|[A-Z])', r'巨蟹座\1\2'],
		[r'(\w+) Vir *(|[A-Z])', r'室女座\1\2'],
		[r'(\w+) Cet *(|[A-Z])', r'鯨魚座\1\2'],
		[r'(\w+) Aqr *(|[A-Z])', r'寶瓶座\1\2'],
		[r'(\w+) Pic *(|[A-Z])', r'繪架座\1\2'],
		[r'(\w+) Cha *(|[A-Z])', r'蝘蜓座\1\2'],
		[r'(\w+) Tau *(|[A-Z])', r'金牛座\1\2'],
		[r'(\w+) Leo *(|[A-Z])', r'獅子座\1\2'],
		[r'(\w+) Psc *(|[A-Z])', r'雙魚座\1\2'],
		[r'(\w+) Ser *(|[A-Z])', r'巨蛇座\1\2'],
		[r'(\w+) Cae *(|[A-Z])', r'鵰具座\1\2'],
		[r'(\w+) For *(|[A-Z])', r'天爐座\1\2'],
		[r'(\w+) Cen *(|[A-Z])', r'半人馬座\1\2'],
		[r'(\w+) Eri *(|[A-Z])', r'波江座\1\2'],
		[r'(\w+) CrB *(|[A-Z])', r'北冕座\1\2'],
		[r'(\w+) Cep *(|[A-Z])', r'仙王座\1\2'],
		[r'(\w+) Oph *(|[A-Z])', r'蛇夫座\1\2'],
		[r'(\w+) Boo *(|[A-Z])', r'牧夫座\1\2'],
		[r'(\w+) Gem *(|[A-Z])', r'雙子座\1\2'],
		[r'(\w+) Aql *(|[A-Z])', r'天鷹座\1\2'],
		[r'(\w+) Her *(|[A-Z])', r'武仙座\1\2'],
		[r'Oph (\d+)', r'蛇夫座\1'],
		['Formalhaut', '北落師門'],
		['Kapteyn', '卡普坦星'],
		['CoRoT-', '柯洛'],
		[r'GJ (\d+) *(|[A-Z])', r'格利澤\1\2'],
		['alf', 'α'],
		['bet', 'β'],
		['eps', 'ε'],
		['gam', 'γ'],
		['kap', 'κ'],
		['mu', 'μ'],
		['nu', 'ν'],
		['ome', 'ω'],
		['omi', 'ο'],
		['rho', 'ρ'],
		['tau', 'τ'],
		['ups', 'υ'],
		['xi', 'ξ'],
		['HR 810', '時鐘座ι'],
		['HD 99492', '獅子座83B'],
		['HD 27442', '網罟座ε'],
		['HD 19994', '鯨魚座94'],
		['HD 134987', '天秤座23'],
		['HD 216437', '印第安座ρ'],
		['HIP 75458', '天龍座ι'],
		['HD 10697', '雙魚座109'],
		['HD 216435', '天鶴座τ'],
		['HD 16141', '鯨魚座79'],
		['BD-17 63', 'BD-17°63'],
		['HD 11977', '水蛇座η2|水蛇座η<sup>2</sup>'],
		['HD 81688', '天貓座41'],
		['HD 110014', '室女座χ'],
		['KOI-351', '開普勒90'],
		['KOI-55', '開普勒70'],
	]
	for entry in lib:
		string = re.sub(entry[0], entry[1], string)
	return string

# Generates link for exoplanet, checking whether there's a space
def link(host, letter):
	global site
	hostname = translate(host)
	pipe = hostname.find('|')
	if pipe > 0:
		hostname = hostname[:pipe]
	pagename1 = hostname + ' ' + letter
	pagename2 = hostname + letter
	if host != hostname:
		return '[[' + pagename2 + '|' + letter + ']]'
	page2 = pywikibot.Page(site, unicode(pagename2, 'utf_8'))
	if page2.exists():
		return '[[' + pagename2 + '|' + letter + ']]'
	else:
		return '[[' + pagename1 + '|' + letter + ']]' # Defaults to no space

# Converts possibly empty strings to float
def tofloat(string):
	if string == '':
		return ''
	else:
		return float(string)

# Round x to n significant figures
def sigfig(x, n):
	if x == '':
		return ''
	elif x == 0:
		return '0'
	else:
		digits = int(floor(log10(abs(x))))
		if digits >= 3:
			num = int(round(x, n-1-digits))
		else:
			num = round(x, n-1-digits)
		return str(num)

site = pywikibot.Site('zh', 'wikipedia')
# Change following paths accordingly.
table = load_csv('planets_nasa_nocom.csv')
outtext = wikify(table)
outfile = open('output.txt', 'w')
outfile.write(outtext)
outfile.close()
print('Data converted to Wikitable.')