Improving helper types

Author: apetenchea

arangoasync/typings.py

@@ -1692,6 +1692,32 @@ def __init__(self, data: Json) -> None:
     def name(self) -> str:
         return cast(str, self._data["name"])
 
+    @property
+    def is_smart(self) -> bool:
+        """Check if the graph is a smart graph."""
+        return cast(bool, self._data.get("isSmart", False))
+
+    @property
+    def is_satellite(self) -> bool:
+        """Check if the graph is a satellite graph."""
+        return cast(bool, self._data.get("isSatellite", False))
+
+    @property
+    def number_of_shards(self) -> Optional[int]:
+        return cast(Optional[int], self._data.get("numberOfShards"))
+
+    @property
+    def replication_factor(self) -> Optional[int | str]:
+        return cast(Optional[int | str], self._data.get("replicationFactor"))
+
+    @property
+    def min_replication_factor(self) -> Optional[int]:
+        return cast(Optional[int], self._data.get("minReplicationFactor"))
+
+    @property
+    def write_concern(self) -> Optional[int]:
+        return cast(Optional[int], self._data.get("writeConcern"))
+
     @property
     def edge_definitions(self) -> Jsons:
         return cast(Jsons, self._data.get("edgeDefinitions", list()))
@@ -1720,15 +1746,18 @@ class GraphOptions(JsonWrapper):
             Enterprise Edition.
         write_concern (int | None): The write concern for new collections in the
             graph.
+
+    References:
+        - `create-a-graph <https://docs.arangodb.com/stable/develop/http-api/graphs/named-graphs/#create-a-graph>`__
     """  # noqa: E501
 
     def __init__(
         self,
-        number_of_shards: Optional[int],
-        replication_factor: Optional[int | str],
-        satellites: Optional[List[str]],
-        smart_graph_attribute: Optional[str],
-        write_concern: Optional[int],
+        number_of_shards: Optional[int] = None,
+        replication_factor: Optional[int | str] = None,
+        satellites: Optional[List[str]] = None,
+        smart_graph_attribute: Optional[str] = None,
+        write_concern: Optional[int] = None,
     ) -> None:
         data: Json = dict()
         if number_of_shards is not None:
@@ -1762,3 +1791,57 @@ def smart_graph_attribute(self) -> Optional[str]:
     @property
     def write_concern(self) -> Optional[int]:
         return cast(Optional[int], self._data.get("writeConcern"))
+
+
+class VertexCollectionOptions(JsonWrapper):
+    """Special options for vertex collection creation.
+
+    Args:
+        satellites (list): An array of collection names that is used to create
+            SatelliteCollections for a (Disjoint) SmartGraph using
+            SatelliteCollections (Enterprise Edition only). Each array element must
+            be a string and a valid collection name.
+
+    References:
+        - `add-a-vertex-collection <https://docs.arangodb.com/stable/develop/http-api/graphs/named-graphs/#add-a-vertex-collection>`__
+    """  # noqa: E501
+
+    def __init__(
+        self,
+        satellites: Optional[List[str]] = None,
+    ) -> None:
+        data: Json = dict()
+        if satellites is not None:
+            data["satellites"] = satellites
+        super().__init__(data)
+
+    @property
+    def satellites(self) -> Optional[List[str]]:
+        return cast(Optional[List[str]], self._data.get("satellites"))
+
+
+class EdgeDefinitionOptions(JsonWrapper):
+    """Special options for edge definition creation.
+
+    Args:
+        satellites (list): An array of collection names that is used to create
+            SatelliteCollections for a (Disjoint) SmartGraph using
+            SatelliteCollections (Enterprise Edition only). Each array element must
+            be a string and a valid collection name.
+
+    References:
+        - `add-an-edge-definition <https://docs.arangodb.com/stable/develop/http-api/graphs/named-graphs/#add-an-edge-definition>`__
+    """  # noqa: E501
+
+    def __init__(
+        self,
+        satellites: Optional[List[str]] = None,
+    ) -> None:
+        data: Json = dict()
+        if satellites is not None:
+            data["satellites"] = satellites
+        super().__init__(data)
+
+    @property
+    def satellites(self) -> Optional[List[str]]:
+        return cast(Optional[List[str]], self._data.get("satellites"))

tests/test_typings.py

@@ -4,6 +4,7 @@
     CollectionInfo,
     CollectionStatus,
     CollectionType,
+    EdgeDefinitionOptions,
     GraphOptions,
     GraphProperties,
     JsonWrapper,
@@ -17,6 +18,7 @@
     QueryProperties,
     QueryTrackingConfiguration,
     UserInfo,
+    VertexCollectionOptions,
 )
 
 
@@ -368,3 +370,19 @@ def test_GraphOptions():
     assert graph_options.satellites == ["satellite1", "satellite2"]
     assert graph_options.smart_graph_attribute == "region"
     assert graph_options.write_concern == 1
+
+
+def test_VertexCollectionOptions():
+    options = VertexCollectionOptions(
+        satellites=["col1", "col2"],
+    )
+
+    assert options.satellites == ["col1", "col2"]
+
+
+def test_EdgeDefinitionOptions():
+    options = EdgeDefinitionOptions(
+        satellites=["col1", "col2"],
+    )
+
+    assert options.satellites == ["col1", "col2"]