From d07091a490c27a637283496d07c3b9bdcf5c2d90 Mon Sep 17 00:00:00 2001 From: Aadya Chinubhai Date: Mon, 27 Jul 2026 21:41:16 -0700 Subject: [PATCH 1/2] Add functionality to support dict unpacking in csp.node Signed-off-by: Aadya Chinubhai --- csp/impl/wiring/node_parser.py | 48 +++++++++++++++++++++++++ csp/tests/test_engine.py | 65 ++++++++++++++++++++++++++++++++++ csp/tests/test_parsing.py | 8 +++++ 3 files changed, 121 insertions(+) diff --git a/csp/impl/wiring/node_parser.py b/csp/impl/wiring/node_parser.py index f1204bd2c..eed6a5e2a 100644 --- a/csp/impl/wiring/node_parser.py +++ b/csp/impl/wiring/node_parser.py @@ -14,6 +14,19 @@ from csp.impl.wiring.base_parser import BaseParser, CspParseError, _pythonic_depr_warning +def _csp_output_kwargs(outputs_by_name, values): + try: + items = values.items() + except AttributeError: + raise TypeError(f"csp.output(**{values!r}) requires a dict-like value") from None + for k, v in items: + try: + proxy = outputs_by_name[k] + except KeyError: + raise KeyError(f"unrecognized output '{k}'") from None + proxy + v + + class _SingleProxyFuncArgResolver(object): class INVALID_VALUE: pass @@ -73,6 +86,7 @@ class NodeParser(BaseParser): _CSP_ENGINE_START_TIME_FUNC = "_engine_start_time" _CSP_ENGINE_END_TIME_FUNC = "_engine_end_time" _CSP_ENGINE_STATS_FUNC = "_csp_engine_stats" + _CSP_OUTPUT_KWARGS_FUNC = "_csp_output_kwargs" _CSP_STOP_ENGINE_FUNC = "_csp_stop_engine" _CSP_IN_REALTIME_FUNC = "_csp_in_realtime" @@ -83,6 +97,7 @@ class NodeParser(BaseParser): _CSP_STOP_ENGINE_FUNC: _cspimpl._csp_stop_engine, _CSP_ENGINE_STATS_FUNC: _cspimpl._csp_engine_stats, _CSP_IN_REALTIME_FUNC: _cspimpl._csp_in_realtime, + _CSP_OUTPUT_KWARGS_FUNC: _csp_output_kwargs, } _SPECIAL_BLOCKS_METH = {"alarms", "state", "start", "stop", "outputs"} @@ -402,6 +417,10 @@ def _parse_output_or_return(self, node, is_return): node.lineno, ) nodes = [] + 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 +483,20 @@ 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._returned_outputs.update(o.name for o in self._signature._outputs if o.name is not None) + nodes.append( + ast.Call( + func=ast.Name(id=self._CSP_OUTPUT_KWARGS_FUNC, ctx=ast.Load()), + args=[self._build_outputs_by_name_dict(node), arg.value], + keywords=[], + lineno=node.lineno, + end_lineno=node.end_lineno, + ) + ) + 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) @@ -505,6 +538,21 @@ def _parse_output_or_return(self, node, is_return): return res + def _build_outputs_by_name_dict(self, node): + 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, + lineno=node.lineno, + end_lineno=node.end_lineno, + ) + def _parse_output(self, node): return self._parse_output_or_return(node=node, is_return=False) diff --git a/csp/tests/test_engine.py b/csp/tests/test_engine.py index 0206ef6bb..5851754b0 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, "unrecognized output '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, "requires a dict-like value"): + 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 From 6e30e5521ff8f5764d4ba2945bea48bf2d6760b4 Mon Sep 17 00:00:00 2001 From: Aadya Chinubhai Date: Wed, 29 Jul 2026 19:16:34 -0700 Subject: [PATCH 2/2] Make dict creation static Signed-off-by: Aadya Chinubhai --- csp/impl/wiring/node_parser.py | 140 ++++++++++++++++++++++----------- csp/tests/test_engine.py | 4 +- 2 files changed, 97 insertions(+), 47 deletions(-) diff --git a/csp/impl/wiring/node_parser.py b/csp/impl/wiring/node_parser.py index eed6a5e2a..c5ee5e2c9 100644 --- a/csp/impl/wiring/node_parser.py +++ b/csp/impl/wiring/node_parser.py @@ -14,19 +14,6 @@ from csp.impl.wiring.base_parser import BaseParser, CspParseError, _pythonic_depr_warning -def _csp_output_kwargs(outputs_by_name, values): - try: - items = values.items() - except AttributeError: - raise TypeError(f"csp.output(**{values!r}) requires a dict-like value") from None - for k, v in items: - try: - proxy = outputs_by_name[k] - except KeyError: - raise KeyError(f"unrecognized output '{k}'") from None - proxy + v - - class _SingleProxyFuncArgResolver(object): class INVALID_VALUE: pass @@ -86,7 +73,6 @@ class NodeParser(BaseParser): _CSP_ENGINE_START_TIME_FUNC = "_engine_start_time" _CSP_ENGINE_END_TIME_FUNC = "_engine_end_time" _CSP_ENGINE_STATS_FUNC = "_csp_engine_stats" - _CSP_OUTPUT_KWARGS_FUNC = "_csp_output_kwargs" _CSP_STOP_ENGINE_FUNC = "_csp_stop_engine" _CSP_IN_REALTIME_FUNC = "_csp_in_realtime" @@ -97,7 +83,6 @@ class NodeParser(BaseParser): _CSP_STOP_ENGINE_FUNC: _cspimpl._csp_stop_engine, _CSP_ENGINE_STATS_FUNC: _cspimpl._csp_engine_stats, _CSP_IN_REALTIME_FUNC: _cspimpl._csp_in_realtime, - _CSP_OUTPUT_KWARGS_FUNC: _csp_output_kwargs, } _SPECIAL_BLOCKS_METH = {"alarms", "state", "start", "stop", "outputs"} @@ -125,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() @@ -417,6 +404,8 @@ 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) @@ -486,16 +475,78 @@ def _parse_output_or_return(self, node, is_return): 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._returned_outputs.update(o.name for o in self._signature._outputs if o.name is not None) - nodes.append( - ast.Call( - func=ast.Name(id=self._CSP_OUTPUT_KWARGS_FUNC, ctx=ast.Load()), - args=[self._build_outputs_by_name_dict(node), arg.value], - keywords=[], - lineno=node.lineno, - end_lineno=node.end_lineno, + 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) @@ -517,28 +568,27 @@ 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, node): + def _build_outputs_by_name_dict(self): keys = [] values = [] for output in self._signature._outputs: @@ -546,12 +596,7 @@ def _build_outputs_by_name_dict(self, node): continue keys.append(ast.Constant(value=output.name)) values.append(self._ts_outproxy_expr(output.name)) - return ast.Dict( - keys=keys, - values=values, - lineno=node.lineno, - end_lineno=node.end_lineno, - ) + return ast.Dict(keys=keys, values=values) def _parse_output(self, node): return self._parse_output_or_return(node=node, is_return=False) @@ -902,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 5851754b0..f6fb196a3 100644 --- a/csp/tests/test_engine.py +++ b/csp/tests/test_engine.py @@ -369,7 +369,7 @@ def foo(x: ts[bool]) -> csp.Outputs(a=ts[int]): if csp.ticked(x): csp.output(**{"bogus": 1}) - with self.assertRaisesRegex(KeyError, "unrecognized output 'bogus'"): + 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): @@ -379,7 +379,7 @@ def foo(x: ts[bool]) -> csp.Outputs(a=ts[int], b=ts[int]): values = [1, 2] # not dict-like csp.output(**values) - with self.assertRaisesRegex(TypeError, "requires a dict-like value"): + 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):