reimplemented exception-darkening code, added support for Flask

cs50 · dmalan · Nov 1, 2017 · Oct 29, 2017 · Oct 29, 2017 · Oct 29, 2017
commit a706a5eadccbf2f528ed2a7c1f1de56e861ff305
diff --git a/src/cs50/__init__.py b/src/cs50/__init__.py
@@ -1,6 +1,12 @@
 import sys
 
-from .cs50 import *
+from .cs50 import eprint, get_char, get_float, get_int, get_string
+try:
+    from .cs50 import get_long
+except:
+    pass
+
+from . import flask
 
 
 class CustomImporter(object):

diff --git a/src/cs50/cs50.py b/src/cs50/cs50.py
@@ -6,28 +6,8 @@
 
 from distutils.sysconfig import get_python_lib
 from os.path import abspath, join
-from termcolor import cprint
-from traceback import extract_tb, format_list, format_exception_only
-
-
-def excepthook(type, value, tb):
-    """
-    Format traceback, darkening entries from global site-packages directories
-    and user-specific site-packages directory.
-
-    https://stackoverflow.com/a/46071447/5156190
-    """
-    packages = tuple(join(abspath(p), "") for p in sys.path[1:])
-    for entry in extract_tb(tb):
-        fmt = format_list((entry,))
-        if (entry[0].startswith(packages)):
-            cprint("".join(fmt), attrs=["dark"], end="", file=sys.stderr)
-        else:
-            cprint("".join(fmt), end="", file=sys.stderr)
-    cprint("".join(format_exception_only(type, value)), end="")
-
-
-sys.excepthook = excepthook
+from termcolor import colored
+from traceback import extract_tb, format_list, format_exception_only, format_exception
 
 
 class flushfile():
@@ -64,6 +44,31 @@ def eprint(*args, **kwargs):
     print(*args, end=end, file=sys.stderr, sep=sep)
 
 
+def formatException(type, value, tb):
+    """
+    Format traceback, darkening entries from global site-packages directories
+    and user-specific site-packages directory.
+
+    https://stackoverflow.com/a/46071447/5156190
+    """
+
+    # Absolute paths to site-packages
+    packages = tuple(join(abspath(p), "") for p in sys.path[1:])
+
+    # Darken lines referring to files in site-packages
+    lines = []
+    for line in format_exception(type, value, tb):
+        matches = re.search(r"^  File \"([^\"]+)\", line \d+, in .+", line)
+        if matches and matches.group(1).startswith(packages):
+            lines += colored(line, attrs=["dark"])
+        else:
+            lines += line
+    return "".join(lines).rstrip()
+
+
+sys.excepthook = lambda type, value, tb: print(formatException(type, value, tb), file=sys.stderr)
+
+
 def get_char(prompt=None):
     """
     Read a line of text from standard input and return the equivalent char;

diff --git a/src/cs50/flask.py b/src/cs50/flask.py
@@ -0,0 +1,33 @@
+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 0.12 (in case logging changes in 0.13)
+    version = StrictVersion(get_distribution("flask").version)
+    assert version >= StrictVersion("0.12") and version < StrictVersion("0.13")
+
+    # Get default logger
+    import flask.logging
+    f = flask.logging.create_logger
+
+    def create_logger(app):
+        """Wrap default logger"""
+
+        # Create default logger
+        logger = f(app)
+
+        # Reformat default logger's exceptions
+        # https://docs.python.org/3/library/logging.html#logging.Formatter.formatException
+        for handler in logger.handlers:
+            handler.formatter.formatException = lambda exc_info: formatException(*exc_info)
+        return logger
+
+    # Replace default logger
+    flask.logging.create_logger = create_logger
+
+except:
+    pass
diff --git a/tests/flask/application.py b/tests/flask/application.py
@@ -0,0 +1,16 @@
+from flask import Flask, render_template
+
+app = Flask(__name__)
+
+import cs50
+import requests
+@app.route("/")
+def index():
+    def f():
+        res = requests.get("cs50.harvard.edu")
+    f()
+    return render_template("index.html")
+
+@app.route("/foo", methods=["POST"])
+def foo():
+    return ""
diff --git a/tests/flask/requirements.txt b/tests/flask/requirements.txt
@@ -0,0 +1,2 @@
+cs50
+Flask
diff --git a/tests/flask/templates/error.html b/tests/flask/templates/error.html
@@ -0,0 +1,10 @@
+<!DOCTYPE html>
+
+<html>
+    <head>
+        <title>error</title>
+    </head>
+    <body>
+        error
+    </body>
+</html>
diff --git a/tests/flask/templates/index.html b/tests/flask/templates/index.html
@@ -0,0 +1,10 @@
+<!DOCTYPE html>
+
+<html>
+    <head>
+        <title>flask</title>
+    </head>
+    <body>
+        flask
+    </body>
+</html>
diff --git a/tests/sqltests.py → tests/sql.py b/tests/sqltests.py → tests/sql.py
diff --git a/tests/sqlite.py b/tests/sqlite.py
@@ -0,0 +1,4 @@
+from cs50 import SQL
+
+db = SQL("sqlite:///sqlite.db")
+db.execute("SELECT 1")