python Tornado基础restful API处理程序

Posted 2021-05-09

# coding=utf-8
import re
import json
import jsonschema
import tornado.web
import tornado.gen
from utils.json_encoder import CustomJSONEncoder
from logger import gen_logger


class APIHandler(tornado.web.RequestHandler):
    '''
    Base restful API handler.
    '''
    CORS_MAX_AGE = 3600 * 24
    CORS_ALLOWED_ORIGINS = map(re.compile, (
        r'https?://([\w]+\.)?example\.com$',
        r'https?://([\w]+\.)?example\.cn$',
        r'https?://[\d\.:]+$',
    ))
    input_schema = None

    def initialize(self, *args, **kwargs):
        pass

    def prepare(self):
        if self.request.method in ('PUT', 'POST', 'PATCH'):
            self.validate_schema()

    def validate_schema(self):
        if not self.input_schema:
            return

        try:
            jsonschema.validate(self.get_request_data(), self.input_schema)
        except jsonschema.ValidationError as e:
            self.set_status(400)
            self.finish({
                'message': e.message
            })

    def get_request_data(self):
        if hasattr(self, '_request_data'):
            return self._request_data

        if self.request.headers.get('Content-Type', '').lower().startswith(
                'application/json'):
            data = json.loads(self.request.body)
        else:
            gen_logger.warn('Only `application/json` is supported in request.')
            data = None

        self._request_data = data
        return data

    def set_default_headers(self):
        origin = self.request.headers.get('Origin')
        if origin and any(p.match(origin) for p in self.CORS_ALLOWED_ORIGINS):
            self.set_header('Access-Control-Allow-Origin', origin)
            self.set_header('Access-Control-Allow-Credentials', 'true')
            if self.request.method == 'OPTIONS':
                self.set_header('Access-Control-Allow-Headers', 'Content-Type')
                self.set_header('Access-Control-Max-Age', self.CORS_MAX_AGE)

    def write(self, chunk):
        '''Dumps json using custom encoder'''
        if isinstance(chunk, dict):
            # render dict to json or jsonp response
            chunk = json.dumps(chunk, cls=CustomJSONEncoder)

            callback = self.get_argument("callback", None)

            if self.settings.get('enable_jsonp', None) and callback:
                chunk = ''.join([callback, '(', chunk, ')'])
                self.set_header("Content-Type", "application/javascript")
            else:
                self.set_header("Content-Type",
                                "application/json; charset=UTF-8")

        super(APIHandler, self).write(chunk)

    def options(self):
        self.finish()