|
1 | | -import logging |
2 | | - |
3 | | -from distutils.version import StrictVersion |
4 | | -from pkg_resources import get_distribution |
5 | | - |
6 | | -from .cs50 import _formatException |
7 | | - |
8 | | -# Try to monkey-patch Flask, if installed |
9 | | -try: |
10 | | - |
11 | | - # Only patch >= 1.0 |
12 | | - _version = StrictVersion(get_distribution("flask").version) |
13 | | - assert _version >= StrictVersion("1.0") |
14 | | - |
15 | | - # Reformat logger's exceptions |
16 | | - # http://flask.pocoo.org/docs/1.0/logging/ |
17 | | - # https://docs.python.org/3/library/logging.html#logging.Formatter.formatException |
18 | | - try: |
19 | | - import flask.logging |
20 | | - flask.logging.default_handler.formatter.formatException = lambda exc_info: _formatException(*exc_info) |
21 | | - except Exception: |
22 | | - pass |
23 | | - |
24 | | - # Enable logging when Flask is in use, |
25 | | - # monkey-patching own SQL module, which shouldn't need to know about Flask |
26 | | - logging.getLogger("cs50").disabled = True |
27 | | - try: |
28 | | - import flask |
29 | | - from .sql import SQL |
30 | | - except ImportError: |
31 | | - pass |
32 | | - else: |
33 | | - _execute_before = SQL.execute |
34 | | - def _execute_after(*args, **kwargs): |
35 | | - disabled = logging.getLogger("cs50").disabled |
36 | | - if flask.current_app: |
37 | | - logging.getLogger("cs50").disabled = False |
38 | | - try: |
39 | | - return _execute_before(*args, **kwargs) |
40 | | - finally: |
41 | | - logging.getLogger("cs50").disabled = disabled |
42 | | - SQL.execute = _execute_after |
43 | | - |
44 | | - # When behind CS50 IDE's proxy, ensure that flask.redirect doesn't redirect from HTTPS to HTTP |
45 | | - # https://werkzeug.palletsprojects.com/en/0.15.x/middleware/proxy_fix/#module-werkzeug.middleware.proxy_fix |
46 | | - from os import getenv |
47 | | - if getenv("CS50_IDE_TYPE") == "online": |
48 | | - try: |
49 | | - import flask |
50 | | - from werkzeug.middleware.proxy_fix import ProxyFix |
51 | | - _flask_init_before = flask.Flask.__init__ |
52 | | - def _flask_init_after(self, *args, **kwargs): |
53 | | - _flask_init_before(self, *args, **kwargs) |
54 | | - self.wsgi_app = ProxyFix(self.wsgi_app, x_proto=1) |
55 | | - flask.Flask.__init__ = _flask_init_after |
56 | | - except: |
57 | | - pass |
58 | | - |
59 | | -except Exception: |
60 | | - pass |
| 1 | +import os |
| 2 | +import pkgutil |
| 3 | +import sys |
| 4 | + |
| 5 | +def _wrap_flask(f): |
| 6 | + if f is None: |
| 7 | + return |
| 8 | + |
| 9 | + from distutils.version import StrictVersion |
| 10 | + from .cs50 import _formatException |
| 11 | + |
| 12 | + if f.__version__ < StrictVersion("1.0"): |
| 13 | + return |
| 14 | + |
| 15 | + f.logging.default_handler.formatter.formatException = lambda exc_info: _formatException(*exc_info) |
| 16 | + |
| 17 | + if os.getenv("CS50_IDE_TYPE") == "online": |
| 18 | + from werkzeug.middleware.proxy_fix import ProxyFix |
| 19 | + _flask_init_before = f.Flask.__init__ |
| 20 | + def _flask_init_after(self, *args, **kwargs): |
| 21 | + _flask_init_before(self, *args, **kwargs) |
| 22 | + self.wsgi_app = ProxyFix(self.wsgi_app, x_proto=1) |
| 23 | + f.Flask.__init__ = _flask_init_after |
| 24 | + |
| 25 | + |
| 26 | +# Flask was imported before cs50 |
| 27 | +if "flask" in sys.modules: |
| 28 | + _wrap_flask(sys.modules["flask"]) |
| 29 | + |
| 30 | +# Flask wasn't imported |
| 31 | +else: |
| 32 | + flask_loader = pkgutil.get_loader('flask') |
| 33 | + if flask_loader: |
| 34 | + _exec_module_before = flask_loader.exec_module |
| 35 | + |
| 36 | + def _exec_module_after(*args, **kwargs): |
| 37 | + _exec_module_before(*args, **kwargs) |
| 38 | + _wrap_flask(sys.modules["flask"]) |
| 39 | + |
| 40 | + flask_loader.exec_module = _exec_module_after |
0 commit comments