Skip to content

ENH: Add Support for dict unpacking in csp.node #730

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
aadya940 wants to merge 2 commits into
base: main
Choose a base branch
from
Open

ENH: Add Support for dict unpacking in csp.node #730

Show file tree
Hide file tree
Changes from 1 commit
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
d07091a
Add functionality to support dict unpacking in csp.node
aadya940 Jul 28, 2026
6e30e55
Make dict creation static
aadya940 Jul 30, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
48 changes: 48 additions & 0 deletions csp/impl/wiring/node_parser.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -14,6 +14,19 @@
from csp.impl.wiring.base_parser import BaseParser, CspParseError, _pythonic_depr_warning


def _csp_output_kwargs(outputs_by_name, values):
Comment thread
AdamGlustein marked this conversation as resolved.
Outdated
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

@AdamGlustein AdamGlustein Jul 28, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Please add a comment here noting we override the ast.Add operator to output a value, elsewise this code looks like it's dead. You may also need to do a ruff noqa to tell ruff to ignore it



class _SingleProxyFuncArgResolver(object):
class INVALID_VALUE:
pass
Expand Down Expand Up @@ -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"

@AdamGlustein AdamGlustein Jul 28, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Nit: you can remove _CSP_OUTPUT_KWARGS_FUNC from the dicts here and just call _csp_output_kwargs in the ast.Call node directly. These dictionaries are meant for functions defined in the C++ code.


_CSP_STOP_ENGINE_FUNC = "_csp_stop_engine"
_CSP_IN_REALTIME_FUNC = "_csp_in_realtime"
Expand All @@ -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"}
Expand Down Expand Up @@ -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(
Expand Down Expand Up @@ -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(

@robambalu robambalu Jul 28, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

we can try to avoid the function call overhead and generator the for loop AST directly here instead
note you can use @node( debug_print=True ) to dump the generated readable code for debugging

@aadya940 aadya940 Jul 29, 2026

Copy link
Copy Markdown
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Yes, I'm working on it, I saw the dict creation per tick is actually more expensive than the jump to the function call, so hopefully minimizing that too, so the same dict can be reused with different runtime values.

@robambalu robambalu Jul 29, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

not sure i follow, where would there be an extra dict creation? ** would be replaced with a loop on the variable being expanded, no new dict being created. we should add an isinstance check to your point above ( #730 (comment) )

@aadya940 aadya940 Jul 29, 2026

Copy link
Copy Markdown
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

After the inline change it would generate:

while True:
    yield
    if +#inp_0:
        values = {'a': 1, 'b': 2}
        for k, v in values.items():
            {'a': #outp_0, 'b': #outp_1}[k] + v

A dict display is a runtime instruction, Python constructs a fresh dict every time it evaluates { }, so this builds one per iteration. The keys never change though, so instead take the { } outside the while True, right after the first yield, by then PyNode init has already patched the real proxies into the frame locals, so the map captures them once and stays valid for the node's lifetime:

yield
#csp_outmap = {'a': #outp_0, 'b': #outp_1}   # built once with proxies
...
while True:
    yield
    if +#inp_0:
        values = {'a': 1, 'b': 2}
        for #k, #v in values.items():
            #csp_outmap[#k] + #v

Hoping I'm not missing something obvious here!

@robambalu robambalu Jul 29, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Right I forgot that you need to lookup the out proxy dynamically but yes of course you would build the map once outside of the generator loop

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)
Expand Down Expand Up @@ -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)

Expand Down
65 changes: 65 additions & 0 deletions csp/tests/test_engine.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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]:
Expand Down
8 changes: 8 additions & 0 deletions csp/tests/test_parsing.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
Loading