Add plum.overload

.github/workflows/ci.yml

@@ -22,6 +22,9 @@ jobs:
         run: |
           python -m pip install --upgrade pip
           python -m pip install --upgrade --no-cache-dir -e '.[dev]'
+      - name: Test (mypy)
+        run: |
+          python run_mypy_assertions.py
       - name: Test
         env:
           GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}

Makefile

@@ -1,11 +1,12 @@
-.PHONY: docmake docopen docinit docremove docupdate install test clean
+.PHONY: install test
 
 PACKAGE := plum
 
 install:
 	pip install -e '.[dev]'
 
 test:
-	pre-commit run --all-files && sleep 0.2 && \
+	mypy tests/typechecked
+	pre-commit run --all-files && \
 		PRAGMA_VERSION=`python -c "import sys; print('.'.join(map(str, sys.version_info[:2])))"` \
 			pytest tests -v --cov=$(PACKAGE) --cov-report html:cover --cov-report term-missing

docs/_toc.yml

@@ -12,6 +12,7 @@ chapters:
   - file: classes
   - file: keyword_arguments
   - file: comparison
+  - file: integration
 - file: advanced_usage
   sections:
   - file: conversion_promotion

docs/api.rst

@@ -43,6 +43,10 @@ iPython AutoReload
 .. automodule:: plum.autoreload
     :members:
 
+Overload Support
+----------------
+.. autofunction:: plum.overload.dispatch
+
 Other Utilities
 ---------------
 .. automodule:: plum.util

docs/integration.md

@@ -0,0 +1,33 @@
+# Integration with Linters and `mypy`
+
+Plum's integration with linters and `mypy` is unfortunately limited.
+Properly supporting multiple dispatch in these tools is challenging for a [variety of reasons](https://github.com/python/mypy/issues/11727).
+In this section, we collect various patterns in which Plum plays nicely with type checking.
+
+## Overload Support
+
+At the moment, the only know pattern in which Plum produces `mypy`-compliant code uses `typing.overload`.
+
+An example is as follows:
+
+```python
+from plum.overload import dispatch, overload
+
+
+@overload
+def f(x: int) -> int:
+    return x
+
+
+@overload
+def f(x: str) -> str:
+    return x
+
+
+@dispatch
+def f(x):
+    pass
+```
+
+In the above, `plum.overload.overload` is `typing_extensions.overload`.
+For this pattern to work in all Python versions, you must use `typing_extensions.overload`, not `typing.overload`.

plum/overload.py

@@ -0,0 +1,21 @@
+import sys
+from typing import Any, Callable, TypeVar
+
+if sys.version_info >= (3, 11):  # pragma: specific no cover 3.7 3.8 3.9 3.10
+    from typing import get_overloads, overload
+else:  # pragma: specific no cover 3.11
+    from typing_extensions import get_overloads, overload
+
+from .function import Function
+
+__all__ = ["overload", "dispatch"]
+
+T = TypeVar("T", bound=Callable[..., Any])
+
+
+def dispatch(f: Callable) -> Function:
+    """Decorator to register a particular signature."""
+    f_plum = Function(f)
+    for method in get_overloads(f):
+        f_plum.dispatch(method)
+    return f_plum

pyproject.toml

@@ -34,6 +34,8 @@ dev = [
     "build",
     "tox",
     "jupyter-book",
+    "mypy",
+    "typing-extensions; python_version<='3.10'",
 ]
 
 [project.urls]

run_mypy_assertions.py

@@ -0,0 +1,57 @@
+import subprocess
+from pathlib import Path
+
+if __name__ == "__main__":
+    source_dir = Path("tests/typechecked")  # Files that must be validated using `mypy`
+
+    # Run `mypy` and get the output.
+    p = subprocess.Popen(
+        ["mypy", source_dir],
+        stdout=subprocess.PIPE,
+        stderr=subprocess.PIPE,
+    )
+    stdout, stderr = p.communicate()
+    assert stderr == b"", "`stderr` must be empty."
+    print("Output of `mypy`:")
+    print(stdout.decode())
+
+    unvalidated_errors = []  # Errors that were not marked as intended in the source
+
+    for line in stdout.decode().splitlines():
+        # Parse line in the output of `mypy`. If it cannot be parsed, just skip it.
+        try:
+            path, line_number, status, message = line.split(":", 3)
+        except ValueError:
+            continue
+
+        # We only need to validate errors.
+        if status.lower().strip() != "error":
+            continue
+
+        # Get the line in the source that caused the error.
+        with open(path, "r") as f:
+            path_lines = f.read().splitlines()
+        source_line = path_lines[int(line_number) - 1]
+
+        # See if the error was intended.
+        try:
+            code, match = source_line.split("# mypy: E: ", 1)
+            validated = match.lower() in message.lower()
+        except ValueError:
+            validated = False
+
+        # If it wasn't intended, record the error.
+        if not validated:
+            unvalidated_errors.append(line)
+
+    # Return failure if there are any unvalidated errors.
+    if unvalidated_errors:
+        print("These errors were not validated:")
+        for error in unvalidated_errors:
+            print(error)
+        exit(1)
+    else:
+        print("All errors were validated!")
+        exit(0)
+else:
+    print(__name__)

tests/typechecked/test_overload.py

@@ -0,0 +1,26 @@
+import pytest
+
+from plum import NotFoundLookupError
+from plum.overload import dispatch, overload
+
+
+@overload
+def f(x: int) -> int:
+    return x
+
+
+@overload
+def f(x: str) -> str:
+    return x
+
+
+@dispatch
+def f(x):
+    pass
+
+
+def test_overload() -> None:
+    assert f(1) == 1
+    assert f("1") == "1"
+    with pytest.raises(NotFoundLookupError):
+        f(1.0)  # mypy: E: [call-overload]