tryout for pylance

Author: cvanelteren

docs/conf.py

@@ -13,6 +13,7 @@
 
 # Import statements
 import datetime
+import inspect
 import logging
 import os
 import re
@@ -637,5 +638,20 @@ def _replace_snippet(match):
             pass
 
 
+def process_signature(
+    app, what, name, obj, options, signature, return_annotation
+):
+    """Use compact signatures marked by UltraPlot only in generated docs."""
+    marked = getattr(obj, "__ultraplot_doc_signature__", None)
+    if marked is None and inspect.ismethod(obj):
+        marked = getattr(obj.__func__, "__ultraplot_doc_signature__", None)
+    if marked is None and inspect.isclass(obj):
+        marked = getattr(obj.__init__, "__ultraplot_doc_signature__", None)
+    if marked is not None:
+        return marked, return_annotation
+    return signature, return_annotation
+
+
 def setup(app):
     app.connect("autodoc-process-docstring", process_docstring)
+    app.connect("autodoc-process-signature", process_signature)

pyproject.toml

@@ -49,6 +49,9 @@ dynamic = ["version"]
 packages = {find = {exclude=["docs*", "baseline*", "logo*"]}}
 include-package-data = true
 
+[tool.setuptools.package-data]
+ultraplot = ["py.typed"]
+
 [tool.setuptools_scm]
 write_to = "ultraplot/_version.py"
 write_to_template = "__version__ = '{version}'\n"

ultraplot/axes/cartesian.py

@@ -7,7 +7,7 @@
 import functools
 import inspect
 from dataclasses import dataclass, field
-from typing import Any, Dict, Optional, Tuple, Union
+from typing import Any, Callable, Dict, Optional, Tuple, TypeVar, Union, cast
 
 import matplotlib.axis as maxis
 import matplotlib.dates as mdates
@@ -40,6 +40,8 @@
 
 __all__ = ["CartesianAxes"]
 
+_F = TypeVar("_F", bound=Callable[..., Any])
+
 
 # Tuple of date converters
 DATE_CONVERTERS = (mdates.DateConverter,)
@@ -1860,7 +1862,7 @@ def get_tightbbox(self, renderer, *args, **kwargs):
         return super().get_tightbbox(renderer, *args, **kwargs)
 
 
-def _capture_explicit_format_keys(func):
+def _capture_explicit_format_keys(func: _F) -> _F:
     """
     Preserve raw keyword names before Python binds them to the format signature.
     """
@@ -1870,7 +1872,7 @@ def wrapper(self, *args, **kwargs):
         kwargs.setdefault("_explicit_format_keys", set(kwargs))
         return func(self, *args, **kwargs)
 
-    return wrapper
+    return cast(_F, wrapper)
 
 
 # tmp

ultraplot/figure.py

@@ -7,6 +7,7 @@
 import inspect
 import os
 from contextlib import ExitStack
+from typing import Callable, TypeVar, cast
 
 try:
     from typing import Any, Iterable, List, Optional, Tuple, Union
@@ -50,6 +51,8 @@
     "Figure",
 ]
 
+_F = TypeVar("_F", bound=Callable[..., Any])
+
 
 def _any_not_none(*values):
     """Return whether at least one value is not ``None``."""
@@ -695,7 +698,7 @@ def _draw_context():
     return canvas
 
 
-def _clear_border_cache(func):
+def _clear_border_cache(func: _F) -> _F:
     """
     Decorator that clears the border cache after function execution.
     """
@@ -707,7 +710,7 @@ def wrapper(self, *args, **kwargs):
             delattr(self, "_cached_border_axes")
         return result
 
-    return wrapper
+    return cast(_F, wrapper)
 
 
 class Figure(mfigure.Figure):
@@ -3387,28 +3390,28 @@ def add_axes(self, rect, **kwargs):
 
     @docstring._concatenate_inherited
     @docstring._snippet_manager
-    def add_subplot(self, *args, **kwargs):
+    def add_subplot(self, *args, **kwargs) -> paxes.Axes:
         """
         %(figure.subplot)s
         """
         return self._add_subplot(*args, **kwargs)
 
     @docstring._snippet_manager
-    def subplot(self, *args, **kwargs):  # shorthand
+    def subplot(self, *args, **kwargs) -> paxes.Axes:  # shorthand
         """
         %(figure.subplot)s
         """
         return self._add_subplot(*args, **kwargs)
 
     @docstring._snippet_manager
-    def add_subplots(self, *args, **kwargs):
+    def add_subplots(self, *args, **kwargs) -> pgridspec.SubplotGrid:
         """
         %(figure.subplots)s
         """
         return self._add_subplots(*args, **kwargs)
 
     @docstring._snippet_manager
-    def subplots(self, *args, **kwargs):
+    def subplots(self, *args, **kwargs) -> pgridspec.SubplotGrid:
         """
         %(figure.subplots)s
         """

ultraplot/gridspec.py

@@ -9,7 +9,7 @@
 from collections.abc import MutableSequence
 from functools import wraps
 from numbers import Integral
-from typing import List, Optional, Tuple, Union
+from typing import Callable, List, Optional, Tuple, TypeVar, Union, cast, overload
 
 import matplotlib.axes as maxes
 import matplotlib.gridspec as mgridspec
@@ -122,8 +122,23 @@ def _dummy_method(*args):
     return _dummy_method
 
 
-def _apply_to_all(func=None, *, doc_key=None):
-    def decorator(f):
+_F = TypeVar("_F", bound=Callable[..., object])
+
+
+@overload
+def _apply_to_all(func: _F, *, doc_key: Optional[str] = None) -> _F: ...
+
+
+@overload
+def _apply_to_all(
+    func: None = None, *, doc_key: Optional[str] = None
+) -> Callable[[_F], _F]: ...
+
+
+def _apply_to_all(
+    func: Optional[_F] = None, *, doc_key: Optional[str] = None
+) -> Union[_F, Callable[[_F], _F]]:
+    def decorator(f: _F) -> _F:
         @wraps(f)
         def wrapper(self, *args, **kwargs):
             objs = self._apply_command(f.__name__, *args, **kwargs)
@@ -158,7 +173,7 @@ def wrapper(self, *args, **kwargs):
 
             wrapper.__doc__ = doc
 
-        return wrapper
+        return cast(_F, wrapper)
 
     if func is not None:
         return decorator(func)
@@ -2051,7 +2066,7 @@ def _validate_item(self, items, scalar=False):
         return items
 
     @docstring._snippet_manager
-    def format(self, **kwargs):
+    def format(self, **kwargs) -> None:
         """
         Call the ``format`` command for the `~SubplotGrid.figure`
         and every axes in the grid.

ultraplot/internals/docstring.py

@@ -23,43 +23,48 @@
 # ... print(*_iter_doc(uplt))
 import inspect
 import re
+from typing import Any, Callable, TypeVar, cast, overload
 
 from . import ic  # noqa: F401
 
+_F = TypeVar("_F", bound=Callable[..., Any])
+_T = TypeVar("_T")
 
-def _obfuscate_kwargs(func):
+
+def _obfuscate_kwargs(func: _F) -> _F:
     """
-    Obfuscate keyword args.
+    Mark keyword arguments as compact in generated API documentation.
     """
     return _obfuscate_signature(func, lambda **kwargs: None)
 
 
-def _obfuscate_params(func):
+def _obfuscate_params(func: _F) -> _F:
     """
-    Obfuscate all parameters.
+    Mark all parameters as compact in generated API documentation.
     """
     return _obfuscate_signature(func, lambda *args, **kwargs: None)
 
 
-def _obfuscate_signature(func, dummy):
+def _obfuscate_signature(func: _F, dummy: Callable[..., Any]) -> _F:
     """
-    Obfuscate a misleading or incomplete call signature.
-    Instead users should inspect the parameter table.
+    Mark a misleading or incomplete signature as compact in generated docs.
+
+    The callable's actual signature remains available to Python and language
+    servers; Sphinx reads the marker below when rendering API headings.
     """
-    # Obfuscate signature by converting to *args **kwargs. Note this does
-    # not change behavior of function! Copy parameters from a dummy function
-    # because I'm too lazy to figure out inspect.Parameters API
-    # See: https://stackoverflow.com/a/33112180/4970632
-    sig = inspect.signature(func)
-    sig_repl = inspect.signature(dummy)
-    func.__signature__ = sig.replace(parameters=tuple(sig_repl.parameters.values()))
+    # Keep the compact signature available to documentation tooling without
+    # changing the callable's runtime signature. Sphinx uses this marker to
+    # avoid filling API headings with inherited or dynamically routed options.
+    setattr(func, "__ultraplot_doc_signature__", str(inspect.signature(dummy)))
     return func
 
 
-def _concatenate_inherited(func, prepend_summary=False):
+def _concatenate_inherited(
+    func: _F, prepend_summary: bool = False
+) -> _F:
     """
     Concatenate docstrings from a matplotlib axes method with a ultraplot
-    axes method and obfuscate the call signature.
+    axes method and mark its generated-documentation signature as compact.
     """
     import matplotlib.axes as maxes
     import matplotlib.figure as mfigure
@@ -102,7 +107,7 @@ def _concatenate_inherited(func, prepend_summary=False):
 """
 
     # Return docstring
-    # NOTE: Also obfuscate parameters to avoid partial coverage of call signatures
+    # Keep generated API headings compact to avoid showing partial call signatures.
     func.__doc__ = inspect.cleandoc(doc)
     func = _obfuscate_params(func)
     return func
@@ -143,17 +148,24 @@ def __missing__(self, key):
             return dict.__getitem__(self, key)
         raise KeyError(key)
 
-    def __call__(self, obj):
+    @overload
+    def __call__(self, obj: str) -> str: ...
+
+    @overload
+    def __call__(self, obj: _T) -> _T: ...
+
+    def __call__(self, obj: _T | str) -> _T | str:
         """
         Add snippets to the string or object using ``%(name)s`` substitution. Here
         ``%(name)s`` is used rather than ``.format`` to support invalid identifiers.
         """
         if isinstance(obj, str):
             obj %= self  # add snippets to a string
         else:
-            obj.__doc__ = inspect.getdoc(obj)  # also dedents the docstring
-            if obj.__doc__:
-                obj.__doc__ %= self  # insert snippets after dedent
+            documented = cast(Any, obj)
+            documented.__doc__ = inspect.getdoc(documented)  # also dedents the docstring
+            if documented.__doc__:
+                documented.__doc__ %= self  # insert snippets after dedent
         return obj
 
     def __setitem__(self, key, value):

ultraplot/internals/inputs.py

@@ -5,6 +5,7 @@
 
 import functools
 import sys
+from typing import Any, Callable, TypeVar, cast
 
 import numpy as np
 import numpy.ma as ma
@@ -21,6 +22,8 @@
 except ModuleNotFoundError:
     Triangulation = object
 
+_F = TypeVar("_F", bound=Callable[..., Any])
+
 
 # Constants
 BASEMAP_FUNCS = (  # default latlon=True
@@ -289,13 +292,15 @@ def _parse_triangulation_inputs(*args, **kwargs):
     return triangulation, z, args[1:], kwargs
 
 
-def _parse_triangulation_with_preprocess(*keys, keywords=None, allow_extra=True):
+def _parse_triangulation_with_preprocess(
+    *keys, keywords=None, allow_extra=True
+) -> Callable[[_F], _F]:
     """
     Combines _parse_triangulation with _preprocess_or_redirect for backwards compatibility.
     """
 
-    def _decorator(func):
-        def triangulation_wrapper(self, *args, **kwargs):
+    def _decorator(func: _F) -> _F:
+        def triangulation_wrapper(self, *args, **kwargs) -> Any:
             triangulation, z, remaining_args, updated_kwargs = (
                 _parse_triangulation_inputs(*args, **kwargs)
             )
@@ -318,14 +323,14 @@ def _tri_cartopy_default(args, kwargs):
 
         # Finally make sure all other metadata is correct
         functools.update_wrapper(final_wrapper, func)
-        return final_wrapper
+        return cast(_F, final_wrapper)
 
     return _decorator
 
 
 def _preprocess_or_redirect(
     *keys, keywords=None, allow_extra=True, cartopy_default_transform=True
-):
+) -> Callable[[_F], _F]:
     """
     Redirect internal plotting calls to native matplotlib methods. Also convert
     keyword args to positional and pass arguments through 'data' dictionary.
@@ -336,12 +341,12 @@ def _preprocess_or_redirect(
     if isinstance(keywords, str):
         keywords = (keywords,)
 
-    def _decorator(func):
+    def _decorator(func: _F) -> _F:
         name = func.__name__
         from . import _kwargs_to_args
 
         @functools.wraps(func)
-        def _preprocess_or_redirect(self, *args, **kwargs):
+        def _preprocess_or_redirect(self, *args, **kwargs) -> Any:
             if getattr(self, "_internal_call", None):
                 # Redirect internal matplotlib call to native function
                 from ..axes import PlotAxes
@@ -404,7 +409,7 @@ def _preprocess_or_redirect(self, *args, **kwargs):
                 # Call main function
                 return func(self, *args, **kwargs)  # call unbound method
 
-        return _preprocess_or_redirect
+        return cast(_F, _preprocess_or_redirect)
 
     return _decorator