將可選參數傳遞給 Pandas 中的函數

Question

這是我的數據框：

df = pd.DataFrame({'sym': ['msft', 'amd', 'bac', 'tsla'], 'close': [100, 30, 70, 80], 'sector': ['tech', 'tech', 'bank', 'auto'], 'watch_list': [1, 2, 2, 1]})

這是我想使用的函數的簡化形式：

def add_volume(df):
  df['volume'] = [1000, 2000, 3000, 4000]
  return df

我想將兩個可選參數（ sector ， watch_list ）傳遞給我的函數，用於選擇我想要的行，例如，如果我像這樣調用函數

add_volume (df, sector=['tech'], watch_list=[1])

它返回第一行。

如果我只是這樣稱呼它

add_volume(df, watch_list = [1] ）

它返回第一行和最后一行。 我希望這兩個參數都是可選的。

Answer 1

您可以使用*args參數。 這允許您在第一個參數（在您的情況下）之后添加任意數量的參數（可選參數），這是強制性的。

嘗試這個：

import pandas as pd
df = pd.DataFrame({'sym': ['msft', 'amd', 'bac', 'tsla'], 'close': [100, 30, 70, 80], 'sector': ['tech', 'tech', 'bank', 'auto'], 'watch_list': [1, 2, 2, 1]})

def add_volume(df, *args, **kwargs):
    sector = kwargs.get('sector', None)
    watch_list = kwargs.get('watch_list', None)

    df['volume'] = [1000, 2000, 3000, 4000]
    if sector and watch_list:
        return df[(df['sector']==sector) & (df['watch_list']==watch_list)]
    elif sector:
        return df[df['sector']==sector]
    elif watch_list:
        return df[df['watch_list']==watch_list]
    else:
        return df

print(add_volume(df, sector='tech', watch_list=1)) 
print(add_volume(df, watch_list=1)) 
print(add_volume(df, sector='tech'))
print(add_volume(df))

輸出：

   close sector   sym  watch_list  volume
0    100   tech  msft           1    1000

   close sector   sym  watch_list  volume
0    100   tech  msft           1    1000
3     80   auto  tsla           1    4000

   close sector   sym  watch_list  volume
0    100   tech  msft           1    1000
1     30   tech   amd           2    2000

   close sector   sym  watch_list  volume
0    100   tech  msft           1    1000
1     30   tech   amd           2    2000
2     70   bank   bac           2    3000
3     80   auto  tsla           1    4000