Add support for ANSI on Windows (#31)

* Add support for ANSI on Windows

This is a belated replacement for the reverted #24.

Improvements this time around:

Moved the relevant code into upstream colorama, so it plays nicely with
other packages using terminal features on Windows, like tqdm.

This also gives us basic ANSI support on any old Windows installs out
there that still don't have the native ANSI support. I think they're all
EOL already so not a big deal, but it's a free bonus, so why not.

Last time, we had to revert due to a bug in handling "fake" consoles
like Jupyter notebooks. This time, we're using code that's already been
out in the wild for ~1 month with no issues reported, so hopefully there
won't be any need for emergency reverts.

* Install colorama type stubs to make mypy happy

* work around typeshed bug

* Test running in jupyter

* Fix jupyter test

* Windows subprocess requires Path's to be stringified

* jupyter install fails on py36, so restrict to py37+

* Better check for jupyter install

Avoids making mypy complain about missing jupyter stubs

* Make colorama import check more robust

* Restore temporary type: ignore so I can test the rest...

* Another approach to notebook testing

* Remove unnecessary # type: ignore

* Maybe this

* windows py36 doesn't support ansi

* Improve comment

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