How to refer to a pandas data frame object in a function?

Question

I feel like this should be pretty straightforward but I've been getting errors.I want to create a function where I can pass in three different data frames (temp_years_df,co2_emissions_df,energy_use_df) and plot trends on the column 'United States' with different y-labels on the plot:

%pylab inline
import seaborn as sns

def US_trends(df):
  US_plot=df.T['United States'].plot()
  plt.xlabel("Year")
  if df=temp_years_df:
      plt.ylabel("Number of People Affected by Extreme Temperatures")
  elif df=co2_emissions_df:
      plt.ylabel("CO2 Emissions")
  elif df=energy_use_df:
      plt.ylabel("Energy Use")
 return US_plot

US_trends(temp_years_df)

But this is the error I received:

 if df=temp_years_df:
      ^
SyntaxError: invalid syntax

I tried df==temp_years_df and it still doesn't work. Can anyone enlighten me? Thanks!

score 2 · Answer 1 · answered Jan 20 '18 at 18:01

2

Try using pandas built-in function:

df1.equals(df2)

which returns True if both are equal, else False.

For this specific case try:

if df.equals(temp_years_df):

answered Jan 20 '18 at 18:01

skt7

1,197
8
21

score 1 · Answer 2 · answered Jan 20 '18 at 18:09

1

You can actually set attributes to a data frame. See Adding meta-information/metadata to pandas DataFrame for more information.

df.df_name = "ten_years"

Then you can check in your function that

if df.df_name == "ten_years":
    doSomething

answered Jan 20 '18 at 18:09

Tai

7,684
3
29
49

score 0 · Accepted Answer · answered Jan 20 '18 at 18:03

I'm not exactly sure what you're asking for, but if you need to provide an identifier / label to assess whether to run specific logic for your input, dictionaries are one pythonic way:

def US_trends(dfs, label):
    df = dfs[label]
    US_plot=df.T['United States'].plot()
    plt.xlabel("Year")
    if label == 'temp_years':
        plt.ylabel("Number of People Affected by Extreme Temperatures")
    else:
        plt.ylabel("Y-label")
    return US_plot

input_dfs = {'temp_years': temp_years_df}
US_trends(input_dfs, 'temp_years')

No problem, feel free to accept my answer if it was useful. – jpp Jan 20 '18 at 19:19 — jpp, Jan 20 '18 at 19:19

How to refer to a pandas data frame object in a function?

3 Answers3