Improve autocompletion for arguments to objects wrapped with use_signature

altair/utils/core.py

@@ -9,19 +9,29 @@
 import sys
 import traceback
 import warnings
+from typing import Callable, TypeVar, Any
 
 import jsonschema
 import pandas as pd
 import numpy as np
 
 from altair.utils.schemapi import SchemaBase
 
+try:
+    from typing import ParamSpec
+except ImportError:
+    # ParamSpec was introduced in typing in Python 3.10
+    from typing_extensions import ParamSpec
+
 try:
     from pandas.api.types import infer_dtype as _infer_dtype
 except ImportError:
     # Import for pandas < 0.20.0
     from pandas.lib import infer_dtype as _infer_dtype
 
+_V = TypeVar("_V")
+_P = ParamSpec("_P")
+
 
 def infer_dtype(value):
     """Infer the dtype of the value.
@@ -552,10 +562,10 @@ def parse_shorthand(
     return attrs
 
 
-def use_signature(Obj):
+def use_signature(Obj: Callable[_P, Any]):
     """Apply call signature and documentation of Obj to the decorated method"""
 
-    def decorate(f):
+    def decorate(f: Callable[..., _V]) -> Callable[_P, _V]:
         # call-signature of f is exposed via __wrapped__.
         # we want it to mimic Obj.__init__
         f.__wrapped__ = Obj.__init__