119 lines
3.8 KiB
Python
119 lines
3.8 KiB
Python
def visit_function(self, stmt: p.Function) -> None:
|
|
env: Environment = Environment(self.env)
|
|
pos: list[Function.Parameter] = []
|
|
mixed: list[Function.Parameter] = []
|
|
kw: list[Function.Parameter] = []
|
|
|
|
def eval_param_type(param: p.Function.Parameter) -> Type:
|
|
default_type: Optional[Type] = None
|
|
if param.default is not None:
|
|
default_type = self.type_of(param.default)
|
|
|
|
if param.type is not None:
|
|
param_type: Type = self.resolve_type_expr(param.type)
|
|
if default_type is not None:
|
|
if not self.types.is_subtype(default_type, param_type):
|
|
self.reporter.error(
|
|
param.location or stmt.location,
|
|
f"Cannot use default value of type {default_type} for parameter of type {param_type}",
|
|
)
|
|
return param_type
|
|
|
|
if default_type is not None:
|
|
return default_type
|
|
|
|
return UnknownType()
|
|
|
|
position: int = 0
|
|
for param in stmt.params.pos:
|
|
pos.append(
|
|
Function.Parameter(
|
|
pos=position,
|
|
name=param.name,
|
|
type=eval_param_type(param),
|
|
required=param.default is None,
|
|
)
|
|
)
|
|
position += 1
|
|
for param in stmt.params.mixed:
|
|
mixed.append(
|
|
Function.Parameter(
|
|
pos=position,
|
|
name=param.name,
|
|
type=eval_param_type(param),
|
|
required=param.default is None,
|
|
)
|
|
)
|
|
position += 1
|
|
for param in stmt.params.kw:
|
|
kw.append(
|
|
Function.Parameter(
|
|
pos=position,
|
|
name=param.name,
|
|
type=eval_param_type(param),
|
|
required=param.default is None,
|
|
)
|
|
)
|
|
position += 1
|
|
|
|
param_spec: ParamSpec = ParamSpec(
|
|
pos=pos,
|
|
mixed=mixed,
|
|
kw=kw,
|
|
)
|
|
all_params: list[Function.Parameter] = pos + mixed + kw
|
|
for param in all_params:
|
|
env.define(param.name, param.type)
|
|
|
|
returns_hint: Optional[Type] = None
|
|
if stmt.returns is not None:
|
|
returns_hint = self.resolve_type_expr(stmt.returns)
|
|
inside_function: Function = Function(
|
|
params=param_spec,
|
|
returns=returns_hint,
|
|
)
|
|
self.env.define(stmt.name, inside_function)
|
|
|
|
returned: bool = self.process_block(stmt.body, env)
|
|
inferred_return: Type = UnknownType()
|
|
if not returned:
|
|
env.return_types.append(UnitType())
|
|
return_types: list[Type] = self.types.reduce_types(env.return_types)
|
|
if len(return_types) == 1:
|
|
inferred_return = return_types[0]
|
|
elif len(return_types) > 1:
|
|
self.reporter.error(
|
|
stmt.location,
|
|
f"Mixed return types: {return_types}",
|
|
)
|
|
|
|
returns: Type = UnknownType()
|
|
if returns_hint is not None:
|
|
assert stmt.returns is not None
|
|
returns = returns_hint
|
|
if not self.is_subtype(inferred_return, returns):
|
|
self.reporter.error(
|
|
stmt.returns.location,
|
|
f"Return type mismatch, annotated {returns} but returns {inferred_return}",
|
|
)
|
|
else:
|
|
returns = inferred_return
|
|
|
|
function: Type = Function(
|
|
params=param_spec,
|
|
returns=returns,
|
|
)
|
|
generic_params: list[TypeVar] = []
|
|
all_types: list[Type] = [param.type for param in all_params] + [returns]
|
|
for type in all_types:
|
|
if isinstance(type, TypeVar):
|
|
if type not in generic_params:
|
|
generic_params.append(type)
|
|
|
|
if len(generic_params) != 0:
|
|
function = GenericType(
|
|
name=stmt.name,
|
|
params=generic_params,
|
|
body=function,
|
|
)
|
|
self.env.define(stmt.name, function) |