Add support for ANSI on Windows

MANIFEST.in

@@ -1 +1,2 @@
 include LICENSE
+recursive-include wasabi/tests/test-data *.ipynb

requirements.txt

@@ -1,3 +1,6 @@
 pytest
 typing_extensions
-mypy
+mypy
+types-colorama
+nbconvert
+ipykernel

setup.cfg

@@ -12,6 +12,8 @@ long_description_content_type = text/markdown
 zip_safe = true
 include_package_data = true
 python_requires = >=3.6
+install_requires =
+    colorama >= 0.4.6; sys_platform == "win32" and python_version >= "3.7"
 
 [flake8]
 ignore = E203, E266, E501, E731, W503, E741

wasabi/tests/test-data/wasabi-test-notebook.ipynb

@@ -0,0 +1,42 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "eb5586f8",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import sys\n",
+    "import wasabi\n",
+    "\n",
+    "wasabi.msg.warn(\"This is a test. This is only a test.\")\n",
+    "if sys.version_info >= (3, 7):\n",
+    "    assert wasabi.util.supports_ansi()\n",
+    "\n",
+    "print(sys.stdout)"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3 (ipykernel)",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.10.7"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}

wasabi/tests/test_jupyter.py

@@ -0,0 +1,36 @@
+from pathlib import Path
+import subprocess
+import os
+import sys
+
+import wasabi
+
+TEST_DATA = Path(__file__).absolute().parent / "test-data"
+WASABI_DIR = Path(wasabi.__file__).absolute().parent.parent
+
+
+def test_jupyter():
+    # This runs some code in a jupyter notebook environment, but without actually
+    # starting up the notebook UI. Historically we once had a bug that caused crashes
+    # when importing wasabi in a jupyter notebook, because they replace
+    # sys.stdout/stderr with custom objects that aren't "real" files/ttys. So this makes
+    # sure that we can import and use wasabi inside a notebook without crashing.
+    env = dict(os.environ)
+    if "PYTHONPATH" in env:
+        env["PYTHONPATH"] = f"{WASABI_DIR}{os.pathsep}{env['PYTHONPATH']}"
+    else:
+        env["PYTHONPATH"] = str(WASABI_DIR)
+    subprocess.run(
+        [
+            sys.executable,
+            "-m",
+            "nbconvert",
+            str(TEST_DATA / "wasabi-test-notebook.ipynb"),
+            "--execute",
+            "--stdout",
+            "--to",
+            "notebook",
+        ],
+        env=env,
+        check=True,
+    )

wasabi/util.py

@@ -200,17 +200,22 @@ def can_render(string: str) -> bool:
 
 def supports_ansi() -> bool:
     """Returns True if the running system's terminal supports ANSI escape
-    sequences for color, formatting etc. and False otherwise. Inspired by
-    Django's solution – hacky, but an okay approximation.
+    sequences for color, formatting etc. and False otherwise.
 
     RETURNS (bool): Whether the terminal supports ANSI colors.
     """
     if os.getenv(ENV_ANSI_DISABLED):
         return False
-    # See: https://stackoverflow.com/q/7445658/6400719
-    supported_platform = sys.platform != "Pocket PC" and (
-        sys.platform != "win32" or "ANSICON" in os.environ
-    )
-    if not supported_platform:
-        return False
+    # We require colorama on Windows Python 3.7+, but we might be running on Unix, or we
+    # might be running on Windows Python 3.6. In both cases, colorama might be missing,
+    # *or* there might by accident happen to be an install of an old version that
+    # doesn't have just_fix_windows_console. So we need to confirm not just that we can
+    # import colorama, but that we can import just_fix_windows_console.
+    try:
+        from colorama import just_fix_windows_console
+    except ImportError:
+        if sys.platform == "win32" and "ANSICON" not in os.environ:
+            return False
+    else:
+        just_fix_windows_console()
     return True