5

I have the following function that I've made which takes a member's sign-in date and the length of their contract and returns a modified dataframe with an expiration date for each member. The function works as intended, but I wanted to know if there was a cleaner way of writing the nested function month_to_num(ind). I know python does not implement cases, but is there a way to rewrite the all the if/elif/else statements?

def renewal_date(df):

    def month_to_num(ind):
        if (df.loc[ind, "Contract Type"] == "1 Month"):
            return 1
        elif (df.loc[ind, "Contract Type"] == "3 Months Prom" or 
              df.loc[ind, "Contract Type"] == "3 Month"):
            return 3
        elif (df.loc[ind, "Contract Type"] == "4 Month Promo"):
            return 4
        elif (df.loc[ind, "Contract Type"] == "6 Months"):
            return 6
        else:
            return 12

    for z in range(0, len(df)): 
        exp_date = (df.loc[z, "Date-Joined"] + 
                    relativedelta(months=+month_to_num(z)))
        df.set_value(z,"Date-Renewal", exp_date)

    return df
FCo
  • 485
  • 5
  • 17
Lukasz
  • 2,476
  • 10
  • 41
  • 51

2 Answers2

3

This is more suited to the Code Review Stack Exchange, but to preserve functionality, you'd ordinarily use a dictionary:

MONTH_NUMS = {
    '1 Month': 1,
    '3 Months Prom': 3,
    '3 Month': 3,
    '4 Month Promo': 4,
    '6 Months': 6,
}

def month_to_num(ind):
    return MONTH_NUMS.get(df.loc[ind, 'Contract Type'], 12)

That said, there's a bit of a code smell here in how we are mapping strings to numbers this way.

Community
  • 1
  • 1
Dietrich Epp
  • 205,541
  • 37
  • 345
  • 415
2

You can use a dictionnary in coordination with the dict.get() method which return a default value if the key doesn't exists in the dictionnary.

def renewal_date(df):
    def month_to_num(ind):
        contract_types = {'1 Month': 1, '3 Months Prom': 3, '4 Month Promo': 4, '6 Months': 6}

        return contract_types.get(df.loc[ind, "Contract Type"], 12)

    for z in range(0, len(df)):
        exp_date = (df.loc[z, "Date-Joined"] +
                    relativedelta(months=+month_to_num(z)))
        df.set_value(z, "Date-Renewal", exp_date)

    return df
Pierre Barre
  • 2,174
  • 1
  • 11
  • 23