Files
TB-Docs/report/code/python_check_function.py

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)