Python sample script for connecting to a Django Rest Framework based API

Django_Rest_Connection.py

__author__ = 'David Karchmer'
'''
Main fucntionality to manage API calls to a Django based server (with django-rest-framework)
Demonstrates how to set token on header for secure gets/posts
'''
import json
import requests
import logging

DOMAIN_NAME = 'https://example.com'
API_PREFIX = 'api/v1'
logger = logging.getLogger(__name__)

class DjangoRestConnection(object):

    token = None
    domain = DOMAIN_NAME

    def __init__(self, domain=None):
        if domain:
            self.domain = domain

    def _get_url(self, api):
        url = '{0}/{1}/{2}'.format(self.domain, API_PREFIX, api)
        logger.debug('Calling: ' + url)
        return url

    def _get_header(self, use_token):
        if use_token:
            if not self.token:
                raise('No Token')
            authorization_str = 'token %s' % self.token
            headers = {'content-type': 'application/json',
                       'Authorization': authorization_str}
        else:
            headers = {'Content-Type': 'application/json'}
        return headers

    def post(self, api, use_token, data=None):

        url = self._get_url(api)
        headers = self._get_header(use_token)

        if data:
            payload = json.dumps(data)
            r = requests.post(url, data=payload, headers=headers)
        else:
            r = requests.post(url, headers=headers)

        return r


    def get(self, api, use_token):

        url = self._get_url(api)
        headers = self._get_header(use_token)

        r = requests.get(url, headers=headers)

        return r


    def login(self, password, email):
        data = { 'email': email, 'password': password }
        api = 'auth/login'

        r = self.post(api=api, data=data, use_token=False)
        if r.status_code == 200:
            content = json.loads(r.content.decode())
            self.token = content['token']
            self.username = content['username']
            logger.info('Welcome @{0} (token: {1})'.format(self.username, self.token))
            return True
        else:
            logger.error('Login failed: ' + str(r.status_code) + ' ' + r.content.decode())
            return False

    def logout(self):
        api = 'auth/logout'
        r = self.post(api=api, use_token=True)
        if r.status_code == 204:
            logger.info('Goodbye @{0}'.format(self.username))
            self.username = None
            self.token = None
        else:
            logger.error('Logout failed: ' + str(r.status_code) + ' ' + r.content.decode())

if __name__ == '__main__':
    # Test
    # Logger Format
    from logging import StreamHandler, Formatter
    FORMAT = '[%(asctime)-15s] %(levelname)-6s %(message)s'
    DATE_FORMAT = '%d/%b/%Y %H:%M:%S'
    formatter = Formatter(fmt=FORMAT, datefmt=DATE_FORMAT)
    handler = StreamHandler()
    handler.setFormatter(formatter)
    logger.addHandler(handler)
    logger.setLevel(logging.DEBUG)

    c = DjangoRestConnection('http://127.0.0.1:8000')
    c.login(email='[email protected]', password='user1')
    r = c.get(api='group', use_token=True)
    logger.info(r.status_code)
    if r.status_code == 200:
        content = json.loads(r.content.decode())
        logger.debug(str(content))

    c.logout()