Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Make the dispatch function more explicit for type checkers with @overload #95

Open
wants to merge 11 commits into
base: master
Choose a base branch
from
2 changes: 1 addition & 1 deletion check_linter_assertions.py
Original file line number Diff line number Diff line change
Expand Up @@ -158,7 +158,7 @@ def run_linter(linter: str) -> str:
str: `stdout`.
"""
p = subprocess.Popen(
[linter, source_dir],
[sys.executable, "-m", linter, source_dir],
stdout=subprocess.PIPE,
stderr=subprocess.PIPE,
)
Expand Down
18 changes: 16 additions & 2 deletions plum/dispatcher.py
Original file line number Diff line number Diff line change
@@ -1,4 +1,4 @@
from typing import Any, Callable, Dict, Optional, Tuple, TypeVar, Union
from typing import Any, Callable, Dict, Optional, Tuple, TypeVar, Union, overload

from .function import Function
from .overload import get_overloads
Expand All @@ -23,7 +23,21 @@ def __init__(self):
self.functions: Dict[str, Function] = {}
self.classes: Dict[str, Dict[str, Function]] = {}

def __call__(self, method: Optional[T] = None, precedence: int = 0) -> T:
@overload
def __call__(self, method: Callable[[Any], Any], precedence: int = 0) -> Function:
...

@overload
def __call__(
self, method: None, precedence: int = 0
) -> Callable[[Callable[[Any], Any]], Function]:
...

def __call__(
self,
method: Optional[Callable[[Any], Any]] = None,
precedence: int = 0,
) -> Union[Function, Callable[[Callable[[Any], Any]], Function]]:
"""Decorator to register for a particular signature.

Args:
Expand Down
2 changes: 1 addition & 1 deletion plum/function.py
Original file line number Diff line number Diff line change
Expand Up @@ -388,7 +388,7 @@ def _handle_not_found_lookup_error(
raise ex
return method, return_type

def __call__(self, *args, **kw_args):
def __call__(self, *args: object, **kw_args: object):
method, return_type = self._resolve_method_with_cache(args=args)
return _convert(method(*args, **kw_args), return_type)

Expand Down