Using key while sorting values for just one column

Lets say we have a df like below:

df = pd.DataFrame({'A':['y2','x3','z1','z1'],'B':['y2','x3','a2','z1']})

    A   B
0   y2  y2
1   x3  x3
2   z1  a2
3   z1  z1

if we wanted to sort the values on just the numbers in column A, we can do:

df.sort_values(by='A',key=lambda x: x.str[1])

    A   B
3   z1  z1
2   z1  a2
0   y2  y2
1   x3  x3

If we wanted to sort by both columns A and B, but have the key only apply to column A, is there a way to do that?

df.sort_values(by=['A','B'],key=lambda x: x.str[1])

Expected output:

    A   B
2   z1  a2
3   z1  z1
0   y2  y2
1   x3  x3

Read more here: https://stackoverflow.com/questions/66278976/using-key-while-sorting-values-for-just-one-column