update requests.session() header after each request

Question

I'm writing a python client for a RESTful API using requests.Session. The API generates and returns a new CSRFToken after each request. The Devs of the API told me I need to update the session headers like so:

csrf = session.cookies.get('CSRFToken')
session.headers.update({'x-csrf-token': csrf})

How can I bind this functionality at a Session object so that this functionality is automatically called after every request (maybe a nice inheritance sollution might be possible) so that it doesn't have to be called by the user after every request?

Answer 1

Here's a minimal example of what you could do:

import requests

class MySession(requests.Session):
    def __init__(self):
        super().__init__()
        self.headers = dict()
    def get(self, url):
        if (c := self.cookies.get('CSRFToken')):
            self.headers['x-csrf-token'] = c
        return super().get(url, headers=self.headers)


with MySession() as session:
    (r := session.get('https://www.google.com')).raise_for_status()