Add Driver helpers for variable lookup#1605
Conversation
jernejfrank
left a comment
jernejfrank
left a comment
There was a problem hiding this comment.
Thanks for tackling this. I left one thought about exposing the variables directly as a dict that might make sense to do later in another PR if it is needed. Lmk, your thoughts :)
| results = [Variable.from_node(n) for n in all_nodes] | ||
| return results | ||
|
|
||
| @functools.cached_property |
There was a problem hiding this comment.
Any particular reason for this to be cached? IMO, if the cache is required, would be better to have this as an internal variable explicitly instead of introducing hidden state
| def variables(self) -> dict[str, Variable]: | ||
| """Returns all variables in the graph keyed by name.""" | ||
| return { | ||
| node_name: Variable.from_node(node_) for node_name, node_ in self.graph.nodes.items() | ||
| } |
There was a problem hiding this comment.
On second thought, if we are exposing get_graph, that should already give us the names of the nodes and combining it with get_variable makes the API consistent and cleaner. So maybe we remove variables for now and can always add it in case it is really desired.
| :return: Matching HamiltonNode. | ||
| :raises KeyError: If the variable does not exist in this Driver's graph. | ||
| """ | ||
| return self.variables[name] |
There was a problem hiding this comment.
We could do build this straight from self.graph so it still retains the hashmap lookup speed
2 participants
Adds public
Driverhelpers for looking up variables by name and retrieving the publicHamiltonGraphrepresentation without iterating overlist_available_variables().Changes
Driver.variablesmapping keyed by variable name.Driver.get_variable(name)for direct variable lookup.Driver.get_graph()to return the externalHamiltonGraphAPI object.How I tested this
uv run --with pytest pytest tests/test_hamilton_driver.py::test_driver_variable_lookup tests/test_hamilton_driver.py::test_driver_get_graph_returns_hamilton_graphuv run ruff format --check hamilton/driver.py tests/test_hamilton_driver.pyuv run ruff check hamilton/driver.py tests/test_hamilton_driver.pyNotes
Fixes #1138.
Checklist