diff --git a/csp/impl/wiring/node_parser.py b/csp/impl/wiring/node_parser.py index f1204bd2c..c5ee5e2c9 100644 --- a/csp/impl/wiring/node_parser.py +++ b/csp/impl/wiring/node_parser.py @@ -110,6 +110,8 @@ def __init__(self, name, raw_func, func_frame, debug_print=False): self._func_globals_modified.update(self._LOCAL_METHODS) self._gen = None + self._uses_outmap = False + # To catch returning from within a for or while loop, which wouldnt work as it seems self._inner_loop_count = 0 self._returned_outputs = set() @@ -402,6 +404,12 @@ def _parse_output_or_return(self, node, is_return): node.lineno, ) nodes = [] + stmts = [] + + for node_arg in node.args: + if isinstance(node_arg, ast.Starred): + raise CspParseError(f"{func_name} does not support * unpacking", node.lineno) + if len(node.args) == 1: if len(self._signature._outputs) > 1 and self._signature.output(0).name is not None: raise CspParseError( @@ -464,6 +472,82 @@ def _parse_output_or_return(self, node, is_return): self._returned_outputs.add(node.args[0].id) for arg in node.keywords: + if arg.arg is None: + # A **expr unpack, resolved at runtime, can't statically verify + # which outputs it covers, so assume it may cover all of them. + self._uses_outmap = True + self._returned_outputs.update( + o.name for o in self._signature._outputs if o.name is not None + ) + # We'll be generating raw statements to build the dict only once per node. + stmts.append( + # #csp_vals = + ast.Assign( + targets=[ast.Name(id="#csp_vals", ctx=ast.Store())], + value=arg.value, + ) + ) + stmts.append( + # if not isinstance(#csp_vals, dict): raise TypeError(...) + ast.If( + test=ast.UnaryOp( + op=ast.Not(), + operand=ast.Call( + func=ast.Name(id="isinstance", ctx=ast.Load()), + args=[ + ast.Name(id="#csp_vals", ctx=ast.Load()), + ast.Name(id="dict", ctx=ast.Load()), + ], + keywords=[], + ), + ), + body=[ + ast.Raise( + exc=ast.Call( + func=ast.Name(id="TypeError", ctx=ast.Load()), + args=[ast.Constant(f"{func_name} argument after ** must be a dict")], + keywords=[], + ), + cause=None, + ) + ], + orelse=[], + ) + ) + stmts.append( + ast.For( + # for #csp_k, #csp_v in #csp_vals.items(): + target=ast.Tuple( + elts=[ + ast.Name(id="#csp_k", ctx=ast.Store()), + ast.Name(id="#csp_v", ctx=ast.Store()) + ], + ctx=ast.Store(), + ), + iter=ast.Call( + func=ast.Attribute( + value=ast.Name(id="#csp_vals", ctx=ast.Load()), + attr="items", + ctx=ast.Load() + ), + args=[], + keywords=[], + ), + # loop body + body=[ast.Expr( + ast.BinOp( + left=ast.Subscript( + value=ast.Name(id="#csp_outmap", ctx=ast.Load()), + slice=ast.Name(id="#csp_k", ctx=ast.Load()), + ctx=ast.Load(), + ), + op=ast.Add(), + right=ast.Name(id="#csp_v", ctx=ast.Load()), + ) + )], + orelse=[], + )) + continue if self._signature.output(arg.arg, True) is None: raise CspParseError(f"unrecognized output '{arg.arg}'", node.lineno) output = self._signature.output(arg.arg) @@ -484,27 +568,36 @@ def _parse_output_or_return(self, node, is_return): nodes.append(ast.BinOp(left=self._ts_outproxy_expr(arg.arg), op=ast.Add(), right=arg.value)) self._returned_outputs.add(arg.arg) - if len(nodes) == 0: + result = [] + if len(nodes) == 0 and len(stmts) == 0: if not is_return: raise CspParseError("Empty output is not allowed", node.lineno) res = ast.Pass(lineno=node.lineno, end_lineno=node.end_lineno) else: - res = ( - ast.BoolOp(op=ast.Or(), values=nodes, lineno=node.lineno, end_lineno=node.end_lineno) - if len(nodes) > 1 - else nodes[0] - ) + if nodes: + res = ( + ast.BoolOp(op=ast.Or(), values=nodes, lineno=node.lineno, end_lineno=node.end_lineno) + if len(nodes) > 1 + else nodes[0] + ) + result.append(ast.Expr(res, lineno=node.lineno, end_lineno=node.end_lineno)) + result.extend(stmts) if is_return: - if isinstance(res, ast.Pass): - return [res, ast.Continue(lineno=node.lineno, end_lineno=node.end_lineno)] - else: - return [ - ast.Expr(res, lineno=node.lineno, end_lineno=node.end_lineno), - ast.Continue(lineno=node.lineno, end_lineno=node.end_lineno), - ] - + return result + [ast.Continue(lineno=node.lineno, end_lineno=node.end_lineno)] + if stmts: + return result return res + def _build_outputs_by_name_dict(self): + keys = [] + values = [] + for output in self._signature._outputs: + if output.name is None: + continue + keys.append(ast.Constant(value=output.name)) + values.append(self._ts_outproxy_expr(output.name)) + return ast.Dict(keys=keys, values=values) + def _parse_output(self, node): return self._parse_output_or_return(node=node, is_return=False) @@ -854,7 +947,12 @@ def _parse_impl(self): # Yield before start block so we can setup stack frame before executing # However, this initial yield shouldn't be within the try-finally block, since if a node does not start, it's stop() logic should not be invoked # This avoids an issue where one node raises an exception upon start(), and then other nodes execute their stop() without having ever started - start_and_body = [ast.Expr(value=ast.Yield(value=None))] + del_vars + start_and_body + outmap_assign = ( + [ast.Assign(targets=[ast.Name(id="#csp_outmap", ctx=ast.Store())], value=self._build_outputs_by_name_dict())] + if self._uses_outmap + else [] + ) + start_and_body = [ast.Expr(value=ast.Yield(value=None))] + del_vars + outmap_assign + start_and_body newbody = init_block + start_and_body newfuncdef = ast.FunctionDef( diff --git a/csp/tests/test_engine.py b/csp/tests/test_engine.py index 0206ef6bb..f6fb196a3 100644 --- a/csp/tests/test_engine.py +++ b/csp/tests/test_engine.py @@ -332,6 +332,71 @@ def count(x: ts[int]) -> ts[int]: result = csp.run(count, x, starttime=datetime(2020, 2, 7, 9), endtime=timedelta(seconds=10))[0] self.assertEqual([v[1] for v in result], list(x * 2 for x in range(1, 11))) + def test_csp_output_dict_unpack(self): + @csp.node + def foo(x: ts[bool]) -> csp.Outputs(a=ts[int], b=ts[int]): + if csp.ticked(x): + values = {"a": 1, "b": 2} + csp.output(**values) + + result = csp.run(foo, csp.const(True), starttime=datetime(2020, 1, 1), endtime=timedelta(seconds=1)) + self.assertEqual(result["a"][0][1], 1) + self.assertEqual(result["b"][0][1], 2) + + def test_csp_output_dict_unpack_mixed(self): + @csp.node + def foo(x: ts[bool]) -> csp.Outputs(a=ts[int], b=ts[int]): + if csp.ticked(x): + csp.output(a=1, **{"b": 2}) + + result = csp.run(foo, csp.const(True), starttime=datetime(2020, 1, 1), endtime=timedelta(seconds=1)) + self.assertEqual(result["a"][0][1], 1) + self.assertEqual(result["b"][0][1], 2) + + def test_csp_output_dict_unpack_multiple(self): + @csp.node + def foo(x: ts[bool]) -> csp.Outputs(a=ts[int], b=ts[int]): + if csp.ticked(x): + csp.output(**{"a": 1}, **{"b": 2}) + + result = csp.run(foo, csp.const(True), starttime=datetime(2020, 1, 1), endtime=timedelta(seconds=1)) + self.assertEqual(result["a"][0][1], 1) + self.assertEqual(result["b"][0][1], 2) + + def test_csp_output_dict_unpack_unknown_key(self): + @csp.node + def foo(x: ts[bool]) -> csp.Outputs(a=ts[int]): + if csp.ticked(x): + csp.output(**{"bogus": 1}) + + with self.assertRaisesRegex(KeyError, "'bogus'"): + csp.run(foo, csp.const(True), starttime=datetime(2020, 1, 1), endtime=timedelta(seconds=1)) + + def test_csp_output_dict_unpack_non_dict(self): + @csp.node + def foo(x: ts[bool]) -> csp.Outputs(a=ts[int], b=ts[int]): + if csp.ticked(x): + values = [1, 2] # not dict-like + csp.output(**values) + + with self.assertRaisesRegex(TypeError, "argument after \\*\\* must be a dict"): + csp.run(foo, csp.const(True), starttime=datetime(2020, 1, 1), endtime=timedelta(seconds=1)) + + def test_csp_output_dict_unpack_basket(self): + @csp.node + def foo(x: ts[bool]) -> csp.Outputs( + a=ts[int], + b=csp.OutputBasket(Dict[str, ts[int]], shape=["k1", "k2"]), + ): + if csp.ticked(x): + values = {"a": 1, "b": {"k1": 10, "k2": 20}} + csp.output(**values) + + result = csp.run(foo, csp.const(True), starttime=datetime(2020, 1, 1), endtime=timedelta(seconds=1)) + self.assertEqual(result["a"][0][1], 1) + self.assertEqual(result["b[k1]"][0][1], 10) + self.assertEqual(result["b[k2]"][0][1], 20) + def test_single_csp_numpy_output(self): @csp.node def count(x: ts[int]) -> ts[int]: diff --git a/csp/tests/test_parsing.py b/csp/tests/test_parsing.py index d89bfb037..8bc2898c2 100644 --- a/csp/tests/test_parsing.py +++ b/csp/tests/test_parsing.py @@ -357,6 +357,14 @@ def foo(x: ts[int]): def foo(x: ts[int]) -> Outputs(x=ts[bool]): __return__(x[1], 7) + with self.assertRaisesRegex(CspParseError, "csp.output does not support \\* unpacking"): + + @csp.node + def foo(x: ts[int]): + __outputs__(z=ts[bool]) + args = (7,) + csp.output(*args) + with self.assertRaisesRegex(CspParseError, "unrecognized output 'x'"): @csp.node