Assume a role in an AWS account with Python

account.py

import boto3


def assume_role(account_id: str, role_name: str) -> boto3.Session:
    sts = boto3.client('sts')
    response = sts.assume_role(
        RoleArn=f'arn:aws:iam::{account_id}:role/{role_name}',
        RoleSessionName='account_creation',
    )
    return boto3.Session(
        response['Credentials']['AccessKeyId'],
        response['Credentials']['SecretAccessKey'],
        response['Credentials']['SessionToken'],
    )