Dread
Dread

Reputation: 861

Python - return dataframe and list from function

I want to create a function in Python that returns both a list and a data frame. I know how to do this with two seperate functions, but is there a way to return both from a single function?

import pandas as pd
# sample data
data_dict = {'unit':['a','b','c','d'],'salary':[100,200,250,300]}
# create data frame
df = pd.DataFrame(data_dict)
# Function that returns a dataframe
def test_func(df):
    # create list meeting some criteria
    mylist = list(df.loc[(df['salary']<=200),'unit'])
    # create new dataframe based on mylist
    new_df = df[df['unit'].isin(mylist)]
    return new_df
# create dataframe from function
new_df = test_func(df)

The above function returns just the data frame. How do I also return mylist in the same function?

Upvotes: 1

Views: 5043

Answers (2)

Ali Hassan
Ali Hassan

Reputation: 966

just place comma and add anything you want, but remember in function calling must give same return variables.

return new_df, my list

new_df, mylist = test_func(df)

Upvotes: 0

Keyb0ardwarri0r
Keyb0ardwarri0r

Reputation: 237

You can simply return two variables in the function

import pandas as pd
# sample data
data_dict = {'unit':['a','b','c','d'],'salary':[100,200,250,300]}
# create data frame
df = pd.DataFrame(data_dict)
# Function that returns a dataframe
def test_func(df):
    # create list meeting some criteria
    mylist = list(df.loc[(df['salary']<=200),'unit'])
    # create new dataframe based on mylist
    new_df = df[df['unit'].isin(mylist)]
    return new_df, my list

# create dataframe and list from function
new_df, mylist = test_func(df)

Upvotes: 4

Related Questions