diff --git a/opto/trace/nodes.py b/opto/trace/nodes.py index 0f4c85e5..ad935015 100644 --- a/opto/trace/nodes.py +++ b/opto/trace/nodes.py @@ -2211,6 +2211,11 @@ def keys(self): return ops.keys(self) + def get(self, key, default=None): + import opto.trace.operators as ops + + return ops.get(self, node(key), node(default)) + def pop(self, __index=-1): # python does hidden type checks import opto.trace.operators as ops diff --git a/opto/trace/operators.py b/opto/trace/operators.py index e319d070..48c03238 100644 --- a/opto/trace/operators.py +++ b/opto/trace/operators.py @@ -430,6 +430,14 @@ def values(x: Dict): return [k for k in x.values()] +@bundle() +def get(x: Any, key: Any, default: Any = None): + """Return the value for key if key is in x, else default.""" + if not hasattr(x, "get"): + raise AttributeError(f"{type(x)} object has no attribute 'get'.") + return x.get(key, default) + + # dict in-place operators