lazily importing flask

Author: Kareem Zidane

src/cs50/__init__.py

@@ -1,28 +1,27 @@
+import logging
 import os
 import sys
 
-try:
-
-    # Save student's sys.path
-    _path = sys.path[:]
 
-    # In case student has files that shadow packages
-    sys.path = [p for p in sys.path if p not in ("", os.getcwd())]
+# Disable cs50 logger by default
+logging.getLogger("cs50").disabled = True
 
-    # Import cs50_*
-    from .cs50 import get_char, get_float, get_int, get_string
+# In case student has files that shadow packages
+for p in ("", os.getcwd()):
     try:
-        from .cs50 import get_long
-    except ImportError:
+        sys.path.remove(p)
+    except ValueError:
         pass
 
-    # Replace Flask's logger
-    from . import flask
-
-    # Wrap SQLAlchemy
-    from .sql import SQL
+# Import cs50_*
+from .cs50 import get_char, get_float, get_int, get_string
+try:
+    from .cs50 import get_long
+except ImportError:
+    pass
 
-finally:
+# Hook into flask importing
+from . import flask
 
-    # Restore student's sys.path (just in case library raised an exception that caller caught)
-    sys.path = _path
+# Wrap SQLAlchemy
+from .sql import SQL

src/cs50/flask.py

@@ -1,60 +1,36 @@
-import logging
-
-from distutils.version import StrictVersion
-from pkg_resources import get_distribution
-
-from .cs50 import _formatException
-
-# Try to monkey-patch Flask, if installed
-try:
-
-    # Only patch >= 1.0
-    _version = StrictVersion(get_distribution("flask").version)
-    assert _version >= StrictVersion("1.0")
-
-    # Reformat logger's exceptions
-    # http://flask.pocoo.org/docs/1.0/logging/
-    # https://docs.python.org/3/library/logging.html#logging.Formatter.formatException
-    try:
-        import flask.logging
-        flask.logging.default_handler.formatter.formatException = lambda exc_info: _formatException(*exc_info)
-    except Exception:
-        pass
-
-    # Enable logging when Flask is in use,
-    # monkey-patching own SQL module, which shouldn't need to know about Flask
-    logging.getLogger("cs50").disabled = True
-    try:
-        import flask
-        from .sql import SQL
-    except ImportError:
-        pass
-    else:
-        _execute_before = SQL.execute
-        def _execute_after(*args, **kwargs):
-            disabled = logging.getLogger("cs50").disabled
-            if flask.current_app:
-                logging.getLogger("cs50").disabled = False
-            try:
-                return _execute_before(*args, **kwargs)
-            finally:
-                logging.getLogger("cs50").disabled = disabled
-        SQL.execute = _execute_after
-
-    # When behind CS50 IDE's proxy, ensure that flask.redirect doesn't redirect from HTTPS to HTTP
-    # https://werkzeug.palletsprojects.com/en/0.15.x/middleware/proxy_fix/#module-werkzeug.middleware.proxy_fix
-    from os import getenv
-    if getenv("CS50_IDE_TYPE") == "online":
-        try:
-            import flask
-            from werkzeug.middleware.proxy_fix import ProxyFix
-            _flask_init_before = flask.Flask.__init__
-            def _flask_init_after(self, *args, **kwargs):
-                _flask_init_before(self, *args, **kwargs)
-                self.wsgi_app = ProxyFix(self.wsgi_app, x_proto=1)
-            flask.Flask.__init__ = _flask_init_after
-        except:
-            pass
-
-except Exception:
-    pass
+import os
+import pkgutil
+import sys
+
+def _wrap_flask(f):
+    from distutils.version import StrictVersion
+
+    if f.__version__ < StrictVersion("1.0"):
+        return
+
+    from .cs50 import _formatException
+    f.logging.default_handler.formatter.formatException = lambda exc_info: _formatException(*exc_info)
+
+    if os.getenv("CS50_IDE_TYPE") == "online":
+        from werkzeug.middleware.proxy_fix import ProxyFix
+        _flask_init_before = f.Flask.__init__
+        def _flask_init_after(self, *args, **kwargs):
+            _flask_init_before(self, *args, **kwargs)
+            self.wsgi_app = ProxyFix(self.wsgi_app, x_proto=1)
+        f.Flask.__init__ = _flask_init_after
+
+
+# Flask was imported before cs50
+if "flask" in sys.modules:
+    _wrap_flask(sys.modules["flask"])
+
+# Flask wasn't imported
+else:
+    flask_loader = pkgutil.get_loader('flask')
+    _exec_module_before = flask_loader.exec_module
+
+    def _exec_module_after(*args, **kwargs):
+        _exec_module_before(*args, **kwargs)
+        _wrap_flask(sys.modules["flask"])
+
+    flask_loader.exec_module = _exec_module_after

src/cs50/sql.py

@@ -1,3 +1,22 @@
+def _enable_logging(f):
+    import logging
+    import functools
+
+    @functools.wraps(f)
+    def decorator(*args, **kwargs):
+        import flask
+
+        disabled = logging.getLogger("cs50").disabled
+        if flask.current_app:
+            logging.getLogger("cs50").disabled = False
+        try:
+            return f(*args, **kwargs)
+        finally:
+            logging.getLogger("cs50").disabled = disabled
+
+    return decorator
+
+
 class SQL(object):
     """Wrap SQLAlchemy to provide a simple SQL API."""
 
@@ -64,6 +83,7 @@ def connect(dbapi_connection, connection_record):
         finally:
             self._logger.disabled = disabled
 
+    @_enable_logging
     def execute(self, sql, *args, **kwargs):
         """Execute a SQL statement."""