gammatools/cybermods.py

import pandas as pd


class CyberMods(object):
    def __init__(self):
        self.raw_cybermods_table = {
            'Score': list(range(1, 11)),
            'Physical': ['Replacement Limb', 'Replacement Sense Organ', 'Attribute Enhancement',
                         'Additional Limb', 'Additional Sense Organ', 'Attachment Interface',
                         'Energy Emitter', 'Energy Absorption', 'Nano-tech Collective', 'YOU INVENT ONE'],
            'Mental': ['Replacement Brain', 'Transmitter', 'Receiver', 'Energy Source Detection', 'Knowledge Base',
                       'Pacify Android', 'Control Machine', 'Fabrication', 'Savant', 'YOU INVENT ONE']
        }

        self.cybermods_table = pd.DataFrame(self.raw_cybermods_table)

    def get_physical_cybermod(self, score):
        cybermod = self.cybermods_table.loc[self.cybermods_table['Score'] == score, 'Physical']
        return cybermod.iloc[0] if not cybermod.empty else None

    def get_mental_cybermod(self, score):
        cybermod = self.cybermods_table.loc[self.cybermods_table['Score'] == score, 'Mental']
        return cybermod.iloc[0] if not cybermod.empty else None

    def get_table_shape(self):
        return self.cybermods_table.shape

    def get_table_row_count(self):
        return self.cybermods_table.shape[0]

    def get_table_column_count(self):
        return self.cybermods_table.shape[1]