Add functionality to support dict unpacking in csp.node

Author: aadya940

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)
 

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]:

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