Scaling: Resource: remember which names a tree holds instead of re-reading it

pylabrobot/resources/deck.py

@@ -1,6 +1,6 @@
 from __future__ import annotations
 
-from typing import Any, Dict, List, Mapping, Optional, cast
+from typing import Any, List, Mapping, Optional, cast
 
 from pylabrobot.resources.errors import ResourceNotFoundError
 
@@ -10,13 +10,7 @@
 
 
 class Deck(Resource):
-  """Base class for liquid handler decks.
-
-  This class maintains a dictionary of all resources on the deck. The dictionary is keyed by the
-  resource name and is updated when resources are assigned and unassigned from the deck. The point
-  of this dictionary is to allow O(1) naming collision checks as well as the quick lookup of
-  resources by name.
-  """
+  """Base class for liquid handler decks."""
 
   def __init__(
     self,
@@ -39,64 +33,16 @@ def __init__(
       metadata=metadata,
     )
     self.location = origin
-    self._resources: Dict[str, Resource] = {}
-
-    self.register_did_assign_resource_callback(self._register_resource)
-    self.register_did_unassign_resource_callback(self._deregister_resource)
 
   def serialize(self) -> dict:
     """Serialize this deck."""
     super_serialized = super().serialize()
     super_serialized.pop("model", None)  # deck's don't typically have a model
     return super_serialized
 
-  def _check_naming_conflicts(self, resource: Resource):
-    """overwrite for speed"""
-    if self.has_resource(resource.name):
-      raise ValueError(f"Resource '{resource.name}' already assigned to deck")
-
-  def _register_resource(self, resource: Resource):
-    """Recursively assign the given resource and all child resources to the `self._resources`
-    dictionary. This method is called after a resource is assigned to the deck
-    (did_assign_resource_callback).
-
-    Precondition: All child resources must be assignable, see `self._check_name_exists`.
-    """
-
-    for child in resource.children:
-      self._register_resource(child)
-    self._resources[resource.name] = resource
-
-  def _deregister_resource(self, resource: Resource):
-    """Recursively deregisters the given resource and all child resources from the `self._resources`
-    dictionary. This method is called after a resource is unassigned from the deck
-    (did_unassign_resource_callback).
-    """
-
-    if self.has_resource(resource.name):
-      del self._resources[resource.name]
-    for child in resource.children:
-      self._deregister_resource(child)
-
-  def get_resource(self, name: str) -> Resource:
-    """Returns the resource with the given name.
-
-    Raises:
-      ResourceNotFoundError: If the resource is not found.
-    """
-    if name == self.name:
-      return self
-    if not self.has_resource(name):
-      raise ResourceNotFoundError(f"Resource '{name}' not found")
-    return self._resources[name]
-
-  def has_resource(self, name: str) -> bool:
-    """Returns True if the deck has a resource with the given name."""
-    return name in self._resources
-
   def get_all_resources(self) -> List[Resource]:
     """Returns a list of all resources in the deck."""
-    return list(self._resources.values())
+    return self.get_all_children()
 
   def clear(self, include_trash: bool = False):
     """Removes all resources from the deck.

pylabrobot/resources/resource.py

@@ -181,11 +181,16 @@ def __init__(
     self.location: Optional[Coordinate] = None
     self.parent: Optional[Resource] = None
     self.children: List[Resource] = []
+    # Everything in this tree, by name, kept only by its root. `assign_child_resource` hands the
+    # map to the new root and `unassign_child_resource` hands it back, the only two moments a root
+    # changes. `None` elsewhere means the names are tracked above, not that there are none.
+    self._subtree_resources: Optional[Dict[str, Resource]] = {name: self}
 
     self._will_assign_resource_callbacks: List[WillAssignResourceCallback] = []
     self._did_assign_resource_callbacks: List[DidAssignResourceCallback] = []
     self._will_unassign_resource_callbacks: List[WillUnassignResourceCallback] = []
     self._did_unassign_resource_callbacks: List[DidUnassignResourceCallback] = []
+
     self._resource_state_updated_callbacks: List[ResourceDidUpdateState] = []
 
   def get_size_x(self) -> float:
@@ -449,7 +454,8 @@ def assign_child_resource(
 
     # Check for unsupported resource assignment operations
     self._check_assignment(resource=resource, reassign=reassign)
-    self.get_root()._check_naming_conflicts(resource=resource)
+    root = self.get_root()
+    arriving = root._check_naming_conflicts(resource=resource)
 
     # Call "will assign" callbacks
     for callback in self._will_assign_resource_callbacks:
@@ -462,6 +468,11 @@ def assign_child_resource(
     resource.location = location
     self.children.append(resource)
 
+    # What arrived belongs to this tree's root now, and no longer heads a tree of its own, so it
+    # gives up the map it was keeping. Collected by the check above, which walked the same subtree.
+    root._resources().update(arriving)
+    resource._subtree_resources = None
+
     # Register callbacks on the new child resource so that they can be propagated up the tree.
     resource.register_will_assign_resource_callback(self._call_will_assign_resource_callbacks)
     resource.register_did_assign_resource_callback(self._call_did_assign_resource_callbacks)
@@ -604,17 +615,43 @@ def is_in_subtree_of(self, other: Resource) -> bool:
       current = current.parent
     return False
 
-  def _check_naming_conflicts(self, resource: Resource):
-    """Recursively check for naming conflicts in the resource tree."""
-    if resource.name == self.name:
-      raise ValueError(f"Resource with name '{resource.name}' already exists in the tree.")
+  def _resources(self) -> Dict[str, Resource]:
+    """The map of names for this tree, which only its root keeps.
 
-    # check if the name of the resource we are currently checking already exists in this subtree
-    for child in self.children:
-      child._check_naming_conflicts(resource)
-    # check if the name of any of the children of the resource already exists in this subtree
-    for child in resource.children:
-      self._check_naming_conflicts(child)
+    Returns:
+      The root's map of every name at or beneath it.
+
+    Raises:
+      RuntimeError: If the root is not holding one, which means a resource stopped heading a tree
+        without handing its map over.
+    """
+    root = self.get_root()
+    if root._subtree_resources is None:
+      raise RuntimeError(f"root '{root.name}' is not holding a map of names")
+    return root._subtree_resources
+
+  def _check_naming_conflicts(self, resource: Resource) -> Dict[str, Resource]:
+    """Raise if anything in `resource`'s subtree is already named in this one.
+
+    Names identify a resource across the whole tree - `get_resource` finds one by name, and
+    `serialize_all_state` keys state by it - so two resources may not share one.
+
+    Args:
+      resource: The resource arriving, with everything beneath it.
+
+    Returns:
+      What arrived, by name, so the caller does not walk the same subtree again to record it.
+
+    Raises:
+      ValueError: If any name in that subtree is already in this tree.
+    """
+    held = self._resources()
+    arriving: Dict[str, Resource] = {}
+    for res in [resource] + resource.get_all_children():
+      if res.name in held:
+        raise ValueError(f"Resource with name '{res.name}' already exists in the tree.")
+      arriving[res.name] = res
+    return arriving
 
   def unassign_child_resource(self, resource: Resource):
     """Unassign a child resource from this resource.
@@ -638,10 +675,18 @@ def unassign_child_resource(self, resource: Resource):
     # Preserve the pose for the event before unassignment clears it.
     previous_location = coordinate_reference(resource.location)
 
+    # The map goes with it: this tree gives up those names and the subtree heads a tree of its
+    # own again, so it takes them back. Read before the tree changes shape.
+    departing = {res.name: res for res in [resource] + resource.get_all_children()}
+    held = self._resources()
+    for name in departing:
+      held.pop(name, None)
+
     # Update the tree structure
     resource.parent = None
     resource.location = None
     self.children.remove(resource)
+    resource._subtree_resources = departing
 
     # Delete callbacks on the child resource so that they are not propagated up the tree.
     resource.deregister_will_assign_resource_callback(self._call_will_assign_resource_callbacks)
@@ -685,16 +730,29 @@ def get_resource(self, name: str) -> Resource:
       ValueError: If no resource with the given name exists.
     """
 
-    if self.name == name:
-      return self
+    resource = self._resources().get(name)
+    if resource is None:
+      raise ResourceNotFoundError(f"Resource with name '{name}' does not exist.")
+    if not (resource is self or resource.is_in_subtree_of(self)):
+      where = (
+        f"assigned to '{resource.parent.name}'" if resource.parent else "the root of this tree"
+      )
+      raise ResourceNotFoundError(
+        f"'{name}' is not at or beneath '{self.name}'. It is in the same tree, {where}."
+      )
+    return resource
 
-    for child in self.children:
-      try:
-        return child.get_resource(name)
-      except ResourceNotFoundError:
-        pass
+  def has_resource(self, name: str) -> bool:
+    """Whether anything at or beneath this resource carries the given name.
+
+    Args:
+      name: The name to look for.
 
-    raise ResourceNotFoundError(f"Resource with name '{name}' does not exist.")
+    Returns:
+      True when a resource with that name is in this subtree.
+    """
+    resource = self._resources().get(name)
+    return resource is not None and (resource is self or resource.is_in_subtree_of(self))
 
   def find_resources(
     self,

pylabrobot/resources/resource_tests.py

@@ -429,10 +429,20 @@ def test_callbacks_removed_on_unassign(self):
     self.r.assign_child_resource(self.child, location=Coordinate.zero())
     self.child.unassign()
 
-    self.assertEqual(self.child._did_assign_resource_callbacks, [])
-    self.assertEqual(self.child._did_unassign_resource_callbacks, [])
-    self.assertEqual(self.child._will_assign_resource_callbacks, [])
-    self.assertEqual(self.child._will_unassign_resource_callbacks, [])
+    # Its own handlers stay; what must go is the parent's, which is what carried an event up.
+
+    self.assertNotIn(
+      self.r._call_did_assign_resource_callbacks, self.child._did_assign_resource_callbacks
+    )
+    self.assertNotIn(
+      self.r._call_did_unassign_resource_callbacks, self.child._did_unassign_resource_callbacks
+    )
+    self.assertNotIn(
+      self.r._call_will_assign_resource_callbacks, self.child._will_assign_resource_callbacks
+    )
+    self.assertNotIn(
+      self.r._call_will_unassign_resource_callbacks, self.child._will_unassign_resource_callbacks
+    )
 
   def test_did_assign_is_passed_up_the_chain(self):
     mock_function = unittest.mock.Mock()
@@ -1184,3 +1194,73 @@ def test_find_resources_no_criteria_returns_self_and_descendants(self):
     self.assertEqual(deck.find_resources(), [deck, plate, trough, waste, well])
     # Non-recursive: self plus direct children only.
     self.assertEqual(deck.find_resources(recursive=False), [deck, plate, trough, waste])
+
+
+class TestNameIndex(unittest.TestCase):
+  """Names are unique across a tree, and the tree remembers which it holds rather than re-reading
+  itself on every assignment. Anything remembered can go stale, so these check it does not."""
+
+  def block(self, name: str) -> Resource:
+    return Resource(name=name, size_x=10, size_y=10, size_z=10)
+
+  def test_a_duplicate_name_is_refused(self):
+    root = self.block("root")
+    root.assign_child_resource(self.block("a"), location=Coordinate.zero())
+    with self.assertRaises(ValueError):
+      root.assign_child_resource(self.block("a"), location=Coordinate.zero())
+
+  def test_a_duplicate_deep_in_the_arriving_subtree_is_refused(self):
+    root = self.block("root")
+    holder = self.block("holder")
+    holder.assign_child_resource(self.block("buried"), location=Coordinate.zero())
+    root.assign_child_resource(holder, location=Coordinate.zero())
+
+    other = self.block("other")
+    other.assign_child_resource(self.block("buried"), location=Coordinate.zero())
+    with self.assertRaises(ValueError):
+      root.assign_child_resource(other, location=Coordinate.zero())
+
+  def test_unassigning_frees_the_name(self):
+    root = self.block("root")
+    plate = self.block("plate")
+    root.assign_child_resource(plate, location=Coordinate.zero())
+    root.unassign_child_resource(plate)
+    root.assign_child_resource(self.block("plate"), location=Coordinate.zero())
+
+  def test_a_subtree_takes_its_names_with_it(self):
+    """The names beneath a resource leave the tree with it, and arrive in whatever tree takes it."""
+    first, second = self.block("first"), self.block("second")
+    holder = self.block("holder")
+    holder.assign_child_resource(self.block("carried"), location=Coordinate.zero())
+    first.assign_child_resource(holder, location=Coordinate.zero())
+
+    # while it is in the first tree, the second knows nothing of what it carries
+    second.assign_child_resource(self.block("carried"), location=Coordinate.zero())
+
+    first.unassign_child_resource(holder)
+    # and now the name it carries collides with the one already there
+    with self.assertRaises(ValueError):
+      second.assign_child_resource(holder, location=Coordinate.zero())
+
+  def test_moving_between_parents_goes_through_the_old_one(self):
+    """A resource is moved by taking it off one parent and putting it on another, in that order.
+
+    Handing it straight to the new parent is refused, because the name is checked while the old
+    parent still holds it. Long-standing behaviour, unrelated to the index, and worth pinning: it is
+    why a plate changing carriers reaches a subscriber as an unassignment and an assignment.
+    """
+    root = self.block("root")
+    left, right = self.block("left"), self.block("right")
+    root.assign_child_resource(left, location=Coordinate.zero())
+    root.assign_child_resource(right, location=Coordinate.zero())
+
+    plate = self.block("plate")
+    left.assign_child_resource(plate, location=Coordinate.zero())
+
+    with self.assertRaises(ValueError):
+      right.assign_child_resource(plate, location=Coordinate.zero())
+
+    left.unassign_child_resource(plate)
+    right.assign_child_resource(plate, location=Coordinate.zero())
+    self.assertIs(plate.parent, right)
+    self.assertEqual(root.get_resource("plate"), plate)