1# -*- coding: utf-8 -*-
2
3"""
4httpbin.filters
5~~~~~~~~~~~~~~~
6
7This module provides response filter decorators.
8"""
9
10import gzip as gzip2
11import zlib
12
13import brotli as _brotli
14
15from six import BytesIO
16from decimal import Decimal
17from time import time as now
18
19from decorator import decorator
20from flask import Flask, Response
21
22
23app = Flask(__name__)
24
25
26@decorator
27def x_runtime(f, *args, **kwargs):
28    """X-Runtime Flask Response Decorator."""
29
30    _t0 = now()
31    r = f(*args, **kwargs)
32    _t1 = now()
33    r.headers['X-Runtime'] = '{0}s'.format(Decimal(str(_t1 - _t0)))
34
35    return r
36
37
38@decorator
39def gzip(f, *args, **kwargs):
40    """GZip Flask Response Decorator."""
41
42    data = f(*args, **kwargs)
43
44    if isinstance(data, Response):
45        content = data.data
46    else:
47        content = data
48
49    gzip_buffer = BytesIO()
50    gzip_file = gzip2.GzipFile(
51        mode='wb',
52        compresslevel=4,
53        fileobj=gzip_buffer
54    )
55    gzip_file.write(content)
56    gzip_file.close()
57
58    gzip_data = gzip_buffer.getvalue()
59
60    if isinstance(data, Response):
61        data.data = gzip_data
62        data.headers['Content-Encoding'] = 'gzip'
63        data.headers['Content-Length'] = str(len(data.data))
64
65        return data
66
67    return gzip_data
68
69
70@decorator
71def deflate(f, *args, **kwargs):
72    """Deflate Flask Response Decorator."""
73
74    data = f(*args, **kwargs)
75
76    if isinstance(data, Response):
77        content = data.data
78    else:
79        content = data
80
81    deflater = zlib.compressobj()
82    deflated_data = deflater.compress(content)
83    deflated_data += deflater.flush()
84
85    if isinstance(data, Response):
86        data.data = deflated_data
87        data.headers['Content-Encoding'] = 'deflate'
88        data.headers['Content-Length'] = str(len(data.data))
89
90        return data
91
92    return deflated_data
93
94
95@decorator
96def brotli(f, *args, **kwargs):
97    """Brotli Flask Response Decorator"""
98
99    data = f(*args, **kwargs)
100
101    if isinstance(data, Response):
102        content = data.data
103    else:
104        content = data
105
106    deflated_data = _brotli.compress(content)
107
108    if isinstance(data, Response):
109        data.data = deflated_data
110        data.headers['Content-Encoding'] = 'br'
111        data.headers['Content-Length'] = str(len(data.data))
112
113        return data
114
115    return deflated_data
116