Using np.where to create a column

I would like to create number of days based on two variables. If the units are in week, multiply the nums by 7, if by month by 12 and 365 if by year. Am unable to achieve this…any leads??

import pandas as pd
import numpy as np

nums = [2, 3, 4, 1, 4, 2, 4, 11, 1, 1]
units = ['days', 'weeks', 'year', 'month', 'days', 'weeks', 'years', 'months', 'day', 'week']


df = pd.DataFrame({'nums': nums,
                  'units': units})

df

def convert_days(var1,var2,var3):
    df[var3]= np.where((df[var1]=='weeks', df[var2]*7, df[var2]))
    return df


convert_days('units', 'nums', 'test')

>Solution :

You can use replace:

dmap = {'days?': 1, 'weeks?': 7, 'months?': 12, 'years?': 365}
df['days'] = df['nums'] * df['units'].replace(dmap, regex=True)
print(df)

# Output
   nums   units  days
0     2    days     2
1     3   weeks    21
2     4    year  1460
3     1   month    12
4     4    days     4
5     2   weeks    14
6     4   years  1460
7    11  months   132
8     1     day     1
9     1    week     7

Note: the s? is to take into account the plural form (ends with ‘s’ or not)