Nodes magic methods for convenience

1 files changed, 51 insertions, 0 deletions

diff --git a/bsfs/graph/nodes.py b/bsfs/graph/nodes.py
index 18ab30d..85e5fdb 100644
--- a/bsfs/graph/nodes.py
+++ b/bsfs/graph/nodes.py
@@ -93,6 +93,57 @@ class Nodes():
         """Return the store's local schema."""
         return self._backend.schema
 
+    def __add__(self, other: typing.Any) -> 'Nodes':
+        """Concatenate guids. Backend, user, and node type must match."""
+        if not isinstance(other, type(self)):
+            return NotImplemented
+        if self._backend != other._backend:
+            raise ValueError(other)
+        if self._user != other._user:
+            raise ValueError(other)
+        if self.node_type != other.node_type:
+            raise ValueError(other)
+        return Nodes(self._backend, self._user, self.node_type, self._guids | other._guids)
+
+    def __or__(self, other: typing.Any) -> 'Nodes':
+        """Concatenate guids. Backend, user, and node type must match."""
+        return self.__add__(other)
+
+    def __sub__(self, other: typing.Any) -> 'Nodes':
+        """Subtract guids. Backend, user, and node type must match."""
+        if not isinstance(other, type(self)):
+            return NotImplemented
+        if self._backend != other._backend:
+            raise ValueError(other)
+        if self._user != other._user:
+            raise ValueError(other)
+        if self.node_type != other.node_type:
+            raise ValueError(other)
+        return Nodes(self._backend, self._user, self.node_type, self._guids - other._guids)
+
+    def __and__(self, other: typing.Any) -> 'Nodes':
+        """Intersect guids. Backend, user, and node type must match."""
+        if not isinstance(other, type(self)):
+            return NotImplemented
+        if self._backend != other._backend:
+            raise ValueError(other)
+        if self._user != other._user:
+            raise ValueError(other)
+        if self.node_type != other.node_type:
+            raise ValueError(other)
+        return Nodes(self._backend, self._user, self.node_type, self._guids & other._guids)
+
+    def __len__(self) -> int:
+        """Return the number of guids."""
+        return len(self._guids)
+
+    def __iter__(self) -> typing.Iterator['Nodes']:
+        """Iterate over individual guids. Returns `Nodes` instances."""
+        return iter(
+            Nodes(self._backend, self._user, self.node_type, {guid})
+            for guid in self._guids
+            )
+
     def __getattr__(self, name: str):
         try:
             return super().__getattr__(name) # type: ignore [misc] # parent has no getattr