ポケモン名からタイプを引く, pandas と openpyxl を使ってエクセルデータを読み取る

こんなエクセルデータの表 pokemon-type-table.xlsx があったとして、それを読み取る。 そして ポケモン名からポケモンタイプを引けるようにする。

pandas と openpyxl は使えるようになっていることが前提。 必要ならインストールしておく。

$ pip3 install pandas
$ pip3 install openpyxl

エクセルデータを読み込んで DataFrame を print する main.py:

import sys
import pandas as pd

def main(xlsxFilePath):
    df = pd.read_excel(xlsxFilePath, sheet_name='Sheet1', dtype={
        'type':     'str',
        'pokemon1': 'str',
        'pokemon2': 'str',
        'pokemon3': 'str',
        'pokemon4': 'str'})

    print(df)

if __name__ == "__main__":
    main('pokemon-type-table.xlsx')

わざわざ列ごとに dtype を指定する必要もないのだが、念の為 str 指定している。

実行してみる。

$ python3 main.py
       type    pokemon1    pokemon2    pokemon3   pokemon4
0    normal       Eevee     Pidgeot  Jigglypuff        NaN
1  electric     Pikachu      Raichu     Voltorb  Electrode
2     water    Squirtle   Wartortle     Psyduck    Golduck
3      fire  Charmander  Charmeleon      Ponyta        NaN
4      rock     Geodude    Graveler       Golem        NaN

次に、 ポケモン名をすべて取得する関数 pokemons を作成する。

行を上から下に走査するには df.iterrows() を使います。

def pokemons(df):
    pokemon_list = []

    # 一行分が (int, pandas.core.series.Series) のタプルとして取り出される.
    for row in df.iterrows():
        series = row[1]
        pokemon_type_name = series['type']
        for colname in ['pokemon1', 'pokemon2', 'pokemon3', 'pokemon4']:
            if series.notnull()[colname]:
                pokemon_name = series[colname]
                pokemon_list.append(pokemon_name)

    return pokemon_list

ポケモンリストをつくるときに、NaN になっているセルは除外したい。 そのためには、 series.notnull()[colname] を使って、そのセルが NaN になっていないかを確認できる。 あとは if 分岐して、セルが NaN でなかった場合だけ ポケモン名を取得する。

print(pokemons(df)) すると:

['Eevee', 'Pidgeot', 'Jigglypuff', 'Pikachu', 'Raichu', 'Voltorb', 'Electrode', 'Squirtle', 'Wartortle', 'Psyduck', 'Golduck', 'Charmander', 'Charmeleon', 'Ponyta', 'Geodude', 'Graveler', 'Golem']

ポケモンがすべて列挙されました。

次にポケモン名からタイプを引く辞書を生成する関数 pokemondictionary をつくります。

def pokemondictionary(df):
    dic = {}

    # 一行分が (int, pandas.core.series.Series) のタプルとして取り出される.
    for row in df.iterrows():
        series = row[1]
        pokemon_type_name = series['type']

        for colname in ['pokemon1', 'pokemon2', 'pokemon3', 'pokemon4']:
            if series.notnull()[colname]:
                pokemon_name = series[colname]
                dic[pokemon_name] = pokemon_type_name
    return dic

すべてのポケモンのタイプを計算してみます。

for pokemon in pokemons(df):
    print( pokemon + ' => ' + dic[pokemon] )

実行してみます。

$ python3 main.py
Eevee => normal
Pidgeot => normal
Jigglypuff => normal
Pikachu => electric
Raichu => electric
Voltorb => electric
Electrode => electric
Squirtle => water
Wartortle => water
Psyduck => water
Golduck => water
Charmander => fire
Charmeleon => fire
Ponyta => fire
Geodude => rock
Graveler => rock
Golem => rock

うまくいきました。

まとめ

コード全体を載せます。

main.py:

import sys
import pandas as pd

# すべてのポケモン名を収集してリストにして返す.
def pokemons(df):
    pokemon_list = []

    # 一行分が (int, pandas.core.series.Series) のタプルとして取り出される.
    for row in df.iterrows():
        series = row[1]
        pokemon_type_name = series['type']
        for colname in ['pokemon1', 'pokemon2', 'pokemon3', 'pokemon4']:
            if series.notnull()[colname]:
                pokemon_name = series[colname]
                pokemon_list.append(pokemon_name)

    return pokemon_list

# ポケモン名からタイプを引くための辞書をつくる.
def pokemondictionary(df):
    dic = {}

    # 一行分が (int, pandas.core.series.Series) のタプルとして取り出される.
    for row in df.iterrows():
        series = row[1]
        pokemon_type_name = series['type']

        for colname in ['pokemon1', 'pokemon2', 'pokemon3', 'pokemon4']:
            if series.notnull()[colname]:
                pokemon_name = series[colname]
                dic[pokemon_name] = pokemon_type_name
    return dic

def main(xlsxFilePath):
    df = pd.read_excel(xlsxFilePath, sheet_name='Sheet1', dtype={
        'type':     'str',
        'pokemon1': 'str',
        'pokemon2': 'str',
        'pokemon3': 'str',
        'pokemon4': 'str'})

    dic = pokemondictionary(df)

    for pokemon in pokemons(df):
        print( pokemon + ' => ' + dic[pokemon] )

if __name__ == "__main__":
    main('pokemon-type-table.xlsx')

pandas 便利。 pandas.core.series.Series の扱いが結構わかりにく。