Graph Collections

arangoasync/collection.py

@@ -1,4 +1,9 @@
-__all__ = ["Collection", "StandardCollection"]
+__all__ = [
+    "Collection",
+    "EdgeCollection",
+    "StandardCollection",
+    "VertexCollection",
+]
 
 
 from typing import Any, Generic, List, Optional, Sequence, Tuple, TypeVar, cast
@@ -1711,3 +1716,73 @@ def response_handler(
             return self.deserializer.loads_many(resp.raw_body)
 
         return await self._executor.execute(request, response_handler)
+
+
+class VertexCollection(Collection[T, U, V]):
+    """Vertex collection API wrapper.
+
+    Args:
+        executor (ApiExecutor): API executor.
+        name (str): Collection name
+        graph (str): Graph name.
+        doc_serializer (Serializer): Document serializer.
+        doc_deserializer (Deserializer): Document deserializer.
+    """
+
+    def __init__(
+        self,
+        executor: ApiExecutor,
+        graph: str,
+        name: str,
+        doc_serializer: Serializer[T],
+        doc_deserializer: Deserializer[U, V],
+    ) -> None:
+        super().__init__(executor, name, doc_serializer, doc_deserializer)
+        self._graph = graph
+
+    def __repr__(self) -> str:
+        return f"<VertexCollection {self.name}>"
+
+    @property
+    def graph(self) -> str:
+        """Return the graph name.
+
+        Returns:
+            str: Graph name.
+        """
+        return self._graph
+
+
+class EdgeCollection(Collection[T, U, V]):
+    """Edge collection API wrapper.
+
+    Args:
+        executor (ApiExecutor): API executor.
+        name (str): Collection name
+        graph (str): Graph name.
+        doc_serializer (Serializer): Document serializer.
+        doc_deserializer (Deserializer): Document deserializer.
+    """
+
+    def __init__(
+        self,
+        executor: ApiExecutor,
+        graph: str,
+        name: str,
+        doc_serializer: Serializer[T],
+        doc_deserializer: Deserializer[U, V],
+    ) -> None:
+        super().__init__(executor, name, doc_serializer, doc_deserializer)
+        self._graph = graph
+
+    def __repr__(self) -> str:
+        return f"<EdgeCollection {self.name}>"
+
+    @property
+    def graph(self) -> str:
+        """Return the graph name.
+
+        Returns:
+            str: Graph name.
+        """
+        return self._graph

arangoasync/database.py

@@ -88,6 +88,40 @@ class Database:
     def __init__(self, executor: ApiExecutor) -> None:
         self._executor = executor
 
+    def _get_doc_serializer(
+        self,
+        doc_serializer: Optional[Serializer[T]] = None,
+    ) -> Serializer[T]:
+        """Figure out the document serializer, defaulting to `Json`.
+
+        Args:
+            doc_serializer (Serializer | None): Optional serializer.
+
+        Returns:
+            Serializer: Either the passed serializer or the default one.
+        """
+        if doc_serializer is None:
+            return cast(Serializer[T], self.serializer)
+        else:
+            return doc_serializer
+
+    def _get_doc_deserializer(
+        self,
+        doc_deserializer: Optional[Deserializer[U, V]] = None,
+    ) -> Deserializer[U, V]:
+        """Figure out the document deserializer, defaulting to `Json`.
+
+        Args:
+            doc_deserializer (Deserializer | None): Optional deserializer.
+
+        Returns:
+            Deserializer: Either the passed deserializer or the default one.
+        """
+        if doc_deserializer is None:
+            return cast(Deserializer[U, V], self.deserializer)
+        else:
+            return doc_deserializer
+
     @property
     def connection(self) -> Connection:
         """Return the HTTP connection."""
@@ -390,17 +424,11 @@ def collection(
         Returns:
             StandardCollection: Collection API wrapper.
         """
-        if doc_serializer is None:
-            serializer = cast(Serializer[T], self.serializer)
-        else:
-            serializer = doc_serializer
-        if doc_deserializer is None:
-            deserializer = cast(Deserializer[U, V], self.deserializer)
-        else:
-            deserializer = doc_deserializer
-
         return StandardCollection[T, U, V](
-            self._executor, name, serializer, deserializer
+            self._executor,
+            name,
+            self._get_doc_serializer(doc_serializer),
+            self._get_doc_deserializer(doc_deserializer),
         )
 
     async def collections(
@@ -604,16 +632,11 @@ async def create_collection(
         def response_handler(resp: Response) -> StandardCollection[T, U, V]:
             if not resp.is_success:
                 raise CollectionCreateError(resp, request)
-            if doc_serializer is None:
-                serializer = cast(Serializer[T], self.serializer)
-            else:
-                serializer = doc_serializer
-            if doc_deserializer is None:
-                deserializer = cast(Deserializer[U, V], self.deserializer)
-            else:
-                deserializer = doc_deserializer
             return StandardCollection[T, U, V](
-                self._executor, name, serializer, deserializer
+                self._executor,
+                name,
+                self._get_doc_serializer(doc_serializer),
+                self._get_doc_deserializer(doc_deserializer),
             )
 
         return await self._executor.execute(request, response_handler)
@@ -661,16 +684,30 @@ def response_handler(resp: Response) -> bool:
 
         return await self._executor.execute(request, response_handler)
 
-    def graph(self, name: str) -> Graph:
+    def graph(
+        self,
+        name: str,
+        doc_serializer: Optional[Serializer[T]] = None,
+        doc_deserializer: Optional[Deserializer[U, V]] = None,
+    ) -> Graph[T, U, V]:
         """Return the graph API wrapper.
 
         Args:
             name (str): Graph name.
+            doc_serializer (Serializer): Custom document serializer.
+                This will be used only for document operations.
+            doc_deserializer (Deserializer): Custom document deserializer.
+                This will be used only for document operations.
 
         Returns:
             Graph: Graph API wrapper.
         """
-        return Graph(self._executor, name)
+        return Graph[T, U, V](
+            self._executor,
+            name,
+            self._get_doc_serializer(doc_serializer),
+            self._get_doc_deserializer(doc_deserializer),
+        )
 
     async def has_graph(self, name: str) -> Result[bool]:
         """Check if a graph exists in the database.
@@ -679,7 +716,7 @@ async def has_graph(self, name: str) -> Result[bool]:
             name (str): Graph name.
 
         Returns:
-            bool: True if the graph exists, False otherwise.
+            bool: `True` if the graph exists, `False` otherwise.
 
         Raises:
             GraphListError: If the operation fails.
@@ -720,17 +757,23 @@ def response_handler(resp: Response) -> List[GraphProperties]:
     async def create_graph(
         self,
         name: str,
+        doc_serializer: Optional[Serializer[T]] = None,
+        doc_deserializer: Optional[Deserializer[U, V]] = None,
         edge_definitions: Optional[Sequence[Json]] = None,
         is_disjoint: Optional[bool] = None,
         is_smart: Optional[bool] = None,
         options: Optional[GraphOptions | Json] = None,
         orphan_collections: Optional[Sequence[str]] = None,
         wait_for_sync: Optional[bool] = None,
-    ) -> Result[Graph]:
+    ) -> Result[Graph[T, U, V]]:
         """Create a new graph.
 
         Args:
             name (str): Graph name.
+            doc_serializer (Serializer): Custom document serializer.
+                This will be used only for document operations.
+            doc_deserializer (Deserializer): Custom document deserializer.
+                This will be used only for document operations.
             edge_definitions (list | None): List of edge definitions, where each edge
                 definition entry is a dictionary with fields "collection" (name of the
                 edge collection), "from" (list of vertex collection names) and "to"
@@ -782,10 +825,15 @@ async def create_graph(
             params=params,
         )
 
-        def response_handler(resp: Response) -> Graph:
-            if resp.is_success:
-                return Graph(self._executor, name)
-            raise GraphCreateError(resp, request)
+        def response_handler(resp: Response) -> Graph[T, U, V]:
+            if not resp.is_success:
+                raise GraphCreateError(resp, request)
+            return Graph[T, U, V](
+                self._executor,
+                name,
+                self._get_doc_serializer(doc_serializer),
+                self._get_doc_deserializer(doc_deserializer),
+            )
 
         return await self._executor.execute(request, response_handler)
 

arangoasync/exceptions.py

@@ -263,6 +263,30 @@ class DocumentUpdateError(ArangoServerError):
     """Failed to update document."""
 
 
+class EdgeCollectionListError(ArangoServerError):
+    """Failed to retrieve edge collections."""
+
+
+class EdgeDefinitionListError(ArangoServerError):
+    """Failed to retrieve edge definitions."""
+
+
+class EdgeDefinitionCreateError(ArangoServerError):
+    """Failed to create edge definition."""
+
+
+class EdgeDefinitionReplaceError(ArangoServerError):
+    """Failed to replace edge definition."""
+
+
+class EdgeDefinitionDeleteError(ArangoServerError):
+    """Failed to delete edge definition."""
+
+
+class EdgeListError(ArangoServerError):
+    """Failed to retrieve edges coming in and out of a vertex."""
+
+
 class GraphCreateError(ArangoServerError):
     """Failed to create the graph."""
 
@@ -275,6 +299,10 @@ class GraphListError(ArangoServerError):
     """Failed to retrieve graphs."""
 
 
+class GraphPropertiesError(ArangoServerError):
+    """Failed to retrieve graph properties."""
+
+
 class IndexCreateError(ArangoServerError):
     """Failed to create collection index."""
 
@@ -389,3 +417,15 @@ class UserReplaceError(ArangoServerError):
 
 class UserUpdateError(ArangoServerError):
     """Failed to update user."""
+
+
+class VertexCollectionCreateError(ArangoServerError):
+    """Failed to create vertex collection."""
+
+
+class VertexCollectionDeleteError(ArangoServerError):
+    """Failed to delete vertex collection."""
+
+
+class VertexCollectionListError(ArangoServerError):
+    """Failed to retrieve vertex collections."""