class TestGeneratorModel(test_factory()):
fe_type = types.Generator(
gen_func=None,
yield_type=types.int32,
arg_types=[types.int64, types.float32],
state_types=[types.intp, types.intp[::1]],
has_finalizer=False,
)
def get_generator_type(self):
"""
Compute the actual generator type (the generator function's return
type is simply "pyobject").
"""
return types.Generator(
gen_func=self.func_ir.func_id.func,
yield_type=types.pyobject,
arg_types=(types.pyobject, ) * self.func_ir.arg_count,
state_types=(types.pyobject, ) * len(self.geninfo.state_vars),
has_finalizer=True,
)
def run_pass(self, state):
state.func_ir = self._strip_phi_nodes(state.func_ir)
state.func_ir._definitions = build_definitions(state.func_ir.blocks)
# Rerun postprocessor to update metadata
post_proc = postproc.PostProcessor(state.func_ir)
post_proc.run(emit_dels=False)
# Ensure we are not in objectmode generator
if state.func_ir.generator_info is not None and state.typemap is not None:
# Rebuild generator type
# TODO: move this into PostProcessor
gentype = state.return_type
state_vars = state.func_ir.generator_info.state_vars
state_types = [state.typemap[k] for k in state_vars]
state.return_type = types.Generator(
gen_func=gentype.gen_func,
yield_type=gentype.yield_type,
arg_types=gentype.arg_types,
state_types=state_types,
has_finalizer=gentype.has_finalizer,
)
return True
