pokemon_type_graph.py

# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.

# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.

# You should have received a copy of the GNU General Public License
# along with this program.  If not, see <http://www.gnu.org/licenses/>.

#   Draws a graph of Pokemon types and their effextiveness against each other.
#   Suitable for pokemon generations II -> IV
#   Solid lines = attack and defense bonus
#   Dashed lines = just attack bonus
#   Dotted lines = just defense bonus
#   Inspired by this post http://www.kaziemdesigns.com/blog/pokemon-and-graphs-part-1
#   Data stolen from https://en.wikipedia.org/wiki/File:Pokemon_Type_Chart.svg
#   Hex codes stolen from http://pokemondb.net/type

#   Requirements: 
#       Install graphviz http://www.graphviz.org/Download.php
#       Install graphviz python wrapper: https://pypi.python.org/pypi/graphviz


from graphviz import Digraph

#   Triangles:
#       Water, Fire, Grass
#       Grass, Ground, Poison
#       Flying, Fighting, Rock
#       Steel, Rock, Fire
#       Fire, Bug, Grass
#       Fighting, Dark, Psychic

#   Sort of triangles but not really
#       Ground, Electric, Water
#       Flying, Fighting, Ice
#       Rock, Fire, Grass
#       Bug, Grass, Rock
#

#   How this data works:
#       Attack is a list of types this type has an attack bonus to (super effective attacks)
#       Defensse is a list of types this type has a defense bonus to (not very / non-effective when attacked)
#       Background and border colour are for sweet pokemon aesthetics

data = {
    'Normal': {
        'defense': ['Rock', 'Steel', 'Ghost'],
        'background': '#8a8a59',
        'border-color': '#79794e',
    },
    'Fire': {
        'attack': ['Ice', 'Grass', 'Bug', 'Steel'],
        'defense': ['Fairy', 'Fire', 'Bug', 'Grass', 'Ice', 'Steel'],
        'background': '#f08030',
        'border-color': '#b4530d'
    },
    'Water': {
        'attack': ['Fire', 'Rock', 'Ground'],
        'defense': ['Fire', 'Water', 'Ice', 'Steel'],
        'background': '#6890f0',
        'border-color': '#1753e3'
    },
    'Electric': {
        'attack': ['Water', 'Flying'],
        'defense': ['Electric', 'Steel', 'Flying'],
        'background': '#f8d030',
        'border-color': '#c19b07'
    },
    'Grass': {
        'attack': ['Water', 'Rock', 'Ground'],
        'defense': ['Water', 'Electric', 'Grass', 'Ground'],
        'background': '#78c850',
        'border-color': '#4a892b'
    },
    'Ice': {
        'attack': ['Dragon', 'Grass', 'Flying', 'Ground'],
        'defense': ['Ice'],
        'background': '#98d8d8',
        'border-color': '#45b6b6'
    },
    'Fighting': {
        'attack': ['Ice', 'Normal', 'Steel', 'Dark', 'Rock'],
        'defense': ['Bug', 'Rock', 'Dark'],
        'background': '#c03028',
        'border-color': '#82211b'
    },
    'Poison': {
        'attack': ['Grass', 'Fairy'],
        'defense': ['Grass', 'Fighting', 'Poison', 'Bug', 'Fairy'],
        'background': '#a040a0',
        'border-color': '#662966'
    },
    'Ground': {
        'attack': ['Fire', 'Poison', 'Electric', 'Steel', 'Rock'],
        'defense': ['Electric', 'Poison', 'Rock'],
        'background': '#e0c068',
        'border-color': '#aa8623'
    },
    'Flying': {
        'attack': ['Fighting', 'Grass', 'Bug'],
        'defense': ['Ground', 'Grass', 'Fighting', 'Bug'],
        'background': '#a890f0',
        'border-color': '#7762b6'
    },
    'Psychic': {
        'attack': ['Fighting', 'Poison'],
        'defense': ['Fighting', 'Psychic'],
        'background': '#f85888',
        'border-color': '#d60945'
    },
    'Bug': {
        'attack': ['Grass', 'Psychic', 'Dark'],
        'defense': ['Grass', 'Fighting', 'Ground'],
        'background': '#a8b820',
        'border-color': '#616b13'
    },
    'Rock': {
        'attack': ['Fire', 'Ice', 'Flying', 'Bug'],
        'defense': ['Normal', 'Fire', 'Poison', 'Flying'],
        'background': '#b8a038',
        'border-color': '#746523'
    },
    'Ghost': {
        'attack': ['Psychic', 'Ghost'],
        'defense': ['Poison', 'Bug'],
        'background': '#705898',
        'border-color': '#413359'
    },
    'Dragon': {
        'attack': ['Dragon'],
        'defense': ['Fire', 'Water', 'Electric', 'Grass'],
        'background': '#7038f8',
        'border-color': '#3d07c0'
    },
    'Dark': {
        'attack': ['Ghost', 'Psychic'],
        'defense': ['Psychic', 'Ghost', 'Dark'],
        'background': '#705848',
        'border-color': '#362a23'
    },
    'Steel': {
        'attack': ['Ice', 'Rock', 'Fairy'],
        'defense': ['Normal', 'Grass', 'Ice', 'Flying', 'Psychic', 'Bug', 'Rock', 'Dragon', 'Steel', 'Fairy'],
        'background': '#b8b8d0',
        'border-color': '#7a7aa7'
    },
    'Fairy': {
        'attack': ['Dark', 'Dragon', 'Fighting'],
        'defense': ['Fighting', 'Bug', 'Dragon', 'Dark'],
        'background': '#e898e8',
        'border-color': '#d547d5'
    },
}

g = Digraph('Pokemon Types', filename='pokemon_types.gv', engine='sfdp')
# uncomment the next line for cool circles
# g = Digraph('Pokemon Types', filename='pokemon_types.gv', engine='circo')
g.edge_attr.update(penwidth='2')
g.node_attr.update(fontname='Arial')
g.graph_attr.update(overlap='scale',  splines='true', sep='1')
attack_edges = []
defense_edges = []
intersection_edges = []

for node_name, type_data in data.items():
    sg = Digraph(node_name)
    if 'background' in type_data:
        sg.node_attr.update(fontcolor='white', fillcolor=type_data.get('background'), style='filled')
    if 'border-color' in type_data:
        line_color = type_data.get('border-color')
        sg.node_attr.update(color=line_color)
    else:
        line_color = 'black'
    sg.node(node_name)
    g.subgraph(sg)

    attacks = set(type_data.get('attack',[]))
    defenses = set(type_data.get('defense', []))
    only_attacks = list(attacks.difference(defenses))
    only_defenses = list(defenses.difference(attacks))
    intersections = list(attacks.intersection(defenses))

    for attack in only_attacks:
        edge = (node_name, attack, line_color)
        attack_edges.append(edge)
    for defense in only_defenses:
        edge = (node_name, defense, line_color)
        defense_edges.append(edge)
    for intersection in intersections:
        edge = (node_name, intersection, line_color)
        intersection_edges.append(edge)

# uncomment the next 2 if statements for a much more complicated graph
# for edge in defense_edges:
#     edge_source, edge_sink, edge_color = edge
#     g.edge(edge_source, edge_sink, color=edge_color, constraint='false', style='dotted', len='10.0')
# for edge in attack_edges:
#     edge_source, edge_sink, edge_color = edge
#     g.edge(edge_source, edge_sink, color=edge_color, style='dashed', constraint='false')
for edge in intersection_edges:
    edge_source, edge_sink, edge_color = edge
    g.edge(edge_source, edge_sink, color=edge_color, style='bold', splines='curved')


g.view()