Make dict creation static

Signed-off-by: Aadya Chinubhai <aadyachinubhai@gmail.com>

Author: aadya940

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 = <expr>
+                    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,41 +568,35 @@ 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:
             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,
-        )
+        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(

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