assume_role.py

import boto3


def assume_role(role_arn, session_name):
    assumed = client.assume_role(RoleArn=role_arn, RoleSessionName=session_name)
    session = boto3.session.Session(aws_access_key_id=assumed['Credentials']['AccessKeyId'],
                      aws_secret_access_key=assumed['Credentials']['SecretAccessKey'],
                      aws_session_token=assumed['Credentials']['SessionToken'])
    
    return session