diff --git a/midas/checker/preamble.py b/midas/checker/preamble.py index 4323762..b35d4dc 100644 --- a/midas/checker/preamble.py +++ b/midas/checker/preamble.py @@ -3,7 +3,15 @@ from typing import Any, Callable, Optional from midas.checker.environment import Environment from midas.checker.registry import TypesRegistry -from midas.checker.types import Function, GenericType, TopType, Type, TypeVar, UnitType +from midas.checker.types import ( + Function, + GenericType, + OverloadedFunction, + TopType, + Type, + TypeVar, + UnitType, +) @dataclass(frozen=True) @@ -70,6 +78,36 @@ class Preamble(Environment): returns=self._types.get_type("int"), ) + T = TypeVar(name="T", bound=None) + self._def_overloads( + name="max", + py_function=max, + signatures=[ + ( + [Param("arg1", T), Param("arg2", T)], + [], + [], + T, + [T], + ), + ([Param("iterable", self._list_of(T))], [], [], T, [T]), + ], + ) + self._def_overloads( + name="min", + py_function=min, + signatures=[ + ( + [Param("arg1", T), Param("arg2", T)], + [], + [], + T, + [T], + ), + ([Param("iterable", self._list_of(T))], [], [], T, [T]), + ], + ) + def _list_of(self, item_type: Type) -> Type: return self._types.apply_generic(self._types.get_type("list"), [item_type]) @@ -142,5 +180,31 @@ class Preamble(Environment): if py_function is not None: self._python_funcs[name] = py_function + def _def_overloads( + self, + *, + name: str, + signatures: list[ + tuple[list[Param], list[Param], list[Param], Type, list[TypeVar]] + ], + py_function: Optional[Callable[..., Any]] = None, + ): + overloads: list[Type] = [] + for pos, mixed, kw, returns, type_vars in signatures: + overloads.append( + self._make_function( + name=name, + pos=pos, + mixed=mixed, + kw=kw, + returns=returns, + type_vars=type_vars, + ) + ) + function: Type = OverloadedFunction(overloads=overloads) + self.define(name, function) + if py_function is not None: + self._python_funcs[name] = py_function + def get_py_func(self, name: str) -> Optional[Callable[..., Any]]: return self._python_funcs.get(name)