User input date shift in Pandas-CodePudding

I have a dataframe, df, where I would like a user to input an integer that will shift specific dates for a given condition.

Data

         start      m1          m2          de          re          type
0       2018-01-01  2022-01-01  2022-04-01  2022-05-01  2022-06-01  aa
1       2018-01-01  2022-01-01  2022-04-01  2022-05-01  2022-06-01  aa
2       2018-01-01  2022-01-01  2022-04-01  2022-05-01  2022-06-01  bb

Desired

     **User input**: 9

        start       m1          m2          de          re          type
0       2018-01-01  2022-10-01  2023-01-01  2023-02-01  2023-03-01  aa
1       2018-01-01  2022-10-01  2023-01-01  2023-02-01  2023-03-01  aa
2       2018-01-01  2022-10-01  2023-01-01  2022-05-01  2022-06-01  bb

#user input the integer: 9. We now add 9 to months in 'm2' column 4 9 we are now in 1/1/2023.

(04/01/2022 is in column ‘m2’. We add 9 months to this date, which gives us 1/1/2023)

Doing

month_value = df["m2"]  = pd.DateOffset(months=int(input()))


 if(df[type] == 'aa'):                                  #condition applied

         d = {
             'm1': pd.DateOffset(months=3),             #we offset m1, 3 months from m2
             'de': pd.DateOffset(months=1),
             're': pd.DateOffset(months=2),
             }
        s = pd.Series(d).rsub(month_value)
        df.assign(**{**s, 'm2': month_value})
    else:
         e = {
             'm1': pd.DateOffset(months=3),
        
             }
        s1 = pd.Series(e).rsub(month_value)
        df.assign(**{**s1, 'm2': month_value})

I am researching. Any suggestion is appreciated.

CodePudding user response：

You can use a lambda function for that:

import pandas as pd

df = pd.DataFrame({'m1': [pd.Timestamp('2022-10-01'), pd.Timestamp('2022-10-01'), pd.Timestamp('2022-10-01')], 
                   'm2': [pd.Timestamp('2022-04-01'), pd.Timestamp('2022-04-01'), pd.Timestamp('2022-04-01')],  
                   'de': [pd.Timestamp('2023-02-01'), pd.Timestamp('2023-02-01'), pd.Timestamp('2022-05-01')],   
                   're': [pd.Timestamp('2023-03-01'), pd.Timestamp('2023-03-01'), pd.Timestamp('2022-06-01')], 
                   'type': ['aa', 'aa', 'bb']})

df["m2"]  = pd.DateOffset(months=int(input('Offset: ')))

# Shift the dependent columns m1, de, re when type is 'aa'
df[df.type == 'aa'] = df[df.type == 'aa'].assign(m1 = lambda row: row['m2'] - pd.DateOffset(months=3))
df[df.type == 'aa'] = df[df.type == 'aa'].assign(de = lambda row: row['m2']   pd.DateOffset(months=1))
df[df.type == 'aa'] = df[df.type == 'aa'].assign(re = lambda row: row['m2']   pd.DateOffset(months=2))

# Shift only m1 when type is not 'aa'
df[df.type != 'aa'] = df[df.type != 'aa'].assign(m1 = lambda row: row['m2'] - pd.DateOffset(months=3))

print(df)

CodePudding user response：

You can use MonthOffset to add the user input month to required column

df = pd.DataFrame({'m1': [pd.Timestamp('2022-10-01'), pd.Timestamp('2022-10-01'), pd.Timestamp('2022-10-01')], 
                   'm2': [pd.Timestamp('2022-04-01'), pd.Timestamp('2022-04-01'), pd.Timestamp('2022-04-01')], 
                   'type': ['aa', 'aa', 'bb']})    

user_input = 9
df['m2'] = df['m2']   pd.offsets.MonthOffset(user_input)

You get

    m1          m2          type
0   2022-10-01  2023-01-01  aa
1   2022-10-01  2023-01-01  aa
2   2022-10-01  2023-01-01  bb