Search code examples
pythonpandasdataframe

How to access multiple values from same row of a dataframe


I have a dataframe containing information about pokemon, and I want to create a pokemon object from a row of the dataframe. Here is my code:

sets = pd.read_excel("SetsSheet.xlsx", index_col=0, header=0)

def fromset(setindex):
    hp = sets.loc[setindex, "hp"]
    type1 = sets.loc[setindex, "type 1"]
    type2 = sets.loc[setindex, "type 2"]
    move1 = sets.loc[setindex, "move 1"]
    move2 = sets.loc[setindex, "move 2"]
    learnset = [move1, move2]
    return BattlePokemon(name, type1, type2, hp, learnset)

This code works fine, but it repeats itself a lot, so I suspect there's a better way I'm supposed to do this, and I want to know what it is.


Solution

  • You can create list of columns and select by them:

    def fromset(setindex):
        
        cols = ["hp", "type 1","type 2", "move 1", "move 2"]
        hp, type1, type2, move1, move2  = sets.loc[setindex, cols]
        learnset = [move1, move2]
        return BattlePokemon(name, type1, type2, hp, learnset)
    

    Also is possible use unpacking for 2 last values to list with *:

    def fromset(setindex):
        
        cols = ["hp", "type 1","type 2", "move 1", "move 2"]
        hp, type1, type2, *learnset  = sets.loc[setindex, cols]
        return BattlePokemon(name, type1, type2, hp, learnset)