Utilize Python Input() to apply conditionals on a dataset-CodePudding

I have a dataset where, whenever a date value in the mig2 column is entered into the input() prompt, mig1, de and re values will be updated according to these rules:

   mig1 is 5 months from the date entered in the input() prompt
   de   is 3 months from the date entered in the input() prompt
   re   is 1 month from the date entered in the input() prompt
   Where [type] column values == 'aa'

The only data that is changing are the dates, which are essentially getting shifted based upon the user input.

Data

location    type    mig1        de          mig2        re
ny          aa      8/1/2021    10/1/2021   1/1/2022    2/1/2022
ny          aa      8/1/2021    10/1/2021   1/1/2022    2/1/2022
ca          aa      8/1/2021    10/1/2021   1/1/2022    2/1/2022
tx          bb      9/1/2021    11/1/2021   2/1/2022    3/1/2022

Desired

 #Date 8/1/2022 is input, which yields:

    location    type    mig1        de          mig2        re
    ny          aa      03/1/2022   05/1/2022   8/1/2022    7/1/2022
    ny          aa      03/1/2022   05/1/2022   8/1/2022    7/1/2022
    ca          aa      03/1/2022   05/1/2022   8/1/2022    7/1/2022
    tx          bb      03/1/2021   11/1/2021   08/1/2022   3/1/2022

Input prompt will ask user which date value they wish to input. User inputs the date '8/1/2021', which updates the remaining column date values. based on rules above.

Doing

datevalue = pd.to_datetime(input("Enter shift: "))


d = {
    'mig1': pd.DateOffset(months=5),
    'de': pd.DateOffset(months=3),
    're': pd.DateOffset(months=1),
}

   e = {
        'mig1': pd.DateOffset(months=5),
        
    }


s = pd.Series(d).rsub(datevalue)
s1 = pd.Series(e).rsub(datevalue)
df.assign(**{**s,s1, 'mig2': datevalue})


Run d if type == 'aa'
Run e if type = 'bb'

Any suggestion is appreciated - Above script works well, but looking for suggestion on incorporating the conditionals

CodePudding user response：

Rather than to use a conditional statement, it's probably better to build a dataframe to update yours. For that, I slightly modified your input:

data = {
    'aa': {'mig1': pd.DateOffset(months=5),
           'de': pd.DateOffset(months=3),
           're': pd.DateOffset(months=1),
           'mig2': pd.DateOffset(0)},
    'bb': {'mig1': pd.DateOffset(months=5),
           'mig2': pd.DateOffset(0)}
}

Now, compute the new values with datevalue, format date to be merged with the original dataframe and create the new dataframe:

data = {typ: {col: (datevalue - offset).strftime('%-m/%-d/%Y')
                 for col, offset in cols.items()}
           for typ, cols in data.items()}

df1 = pd.DataFrame(data).T.loc[df['type']].set_index(df.index)
print(df1)

# Output
  type      mig1        de        re      mig2
0   aa  3/1/2022  5/1/2022  7/1/2022  8/1/2022
1   aa  3/1/2022  5/1/2022  7/1/2022  8/1/2022
2   aa  3/1/2022  5/1/2022  7/1/2022  8/1/2022
3   bb  3/1/2022       NaN       NaN  8/1/2022

Finally update your dataframe:

df.update(df1)
print(df)

# Output
  location type      mig1         de      mig2        re
0       ny   aa  3/1/2022   5/1/2022  8/1/2022  7/1/2022
1       ny   aa  3/1/2022   5/1/2022  8/1/2022  7/1/2022
2       ca   aa  3/1/2022   5/1/2022  8/1/2022  7/1/2022
3       tx   bb  3/1/2022  11/1/2021  8/1/2022  3/1/2022