Allow list of str for algo categories as well

Author: FlorentinD

graphdatascience/session/dedicated_sessions.py

@@ -29,13 +29,17 @@ def estimate(
         self,
         node_count: int,
         relationship_count: int,
-        algorithm_categories: list[AlgorithmCategory] | None = None,
+        algorithm_categories: list[AlgorithmCategory] | list[str] | None = None,
         node_label_count: int = 0,
         node_property_count: int = 0,
         relationship_property_count: int = 0,
     ) -> SessionMemory:
         if algorithm_categories is None:
             algorithm_categories = []
+        else:
+            algorithm_categories = [
+                AlgorithmCategory(cat) if isinstance(cat, str) else cat for cat in algorithm_categories
+            ]
         estimation = self._aura_api.estimate_size(
             node_count=node_count,
             node_label_count=node_label_count,

graphdatascience/session/gds_sessions.py

@@ -55,8 +55,10 @@ def __init__(self, api_credentials: AuraAPICredentials) -> None:
         """
         Initializes a new instance of the GdsSessions class.
 
-        Args:
-            api_credentials (AuraAPICredentials): The Aura API credentials used for establishing a connection.
+        Parameters
+        ----------
+        api_credentials
+            The Aura API credentials used for establishing a connection.
         """
         aura_env = os.environ.get("AURA_ENV")
         aura_api = AuraApi(
@@ -71,23 +73,34 @@ def estimate(
         self,
         node_count: int,
         relationship_count: int,
-        algorithm_categories: Optional[list[AlgorithmCategory]] = None,
+        algorithm_categories: list[AlgorithmCategory] | list[str] | None = None,
         node_label_count: int = 0,
         node_property_count: int = 0,
         relationship_property_count: int = 0,
     ) -> SessionMemory:
         """
         Estimates the memory required for a session with the given node and relationship counts.
 
-        Args:
-            node_count (int): The number of nodes.
-            relationship_count (int): The number of relationships.
-            algorithm_categories (Optional[list[AlgorithmCategory]]): The algorithm categories to consider.
-            node_label_count (int): The number of node labels.
-            node_property_count (int): The number of node properties.
-            relationship_property_count (int): The number of relationship properties.
-        Returns:
-            SessionMemory: The estimated memory required for the session.
+        Parameters
+        ----------
+        node_count
+            Number of nodes.
+        relationship_count
+            Number of relationships.
+        algorithm_categories
+            The algorithm categories to consider.
+        node_label_count
+            Number of node labels.
+        node_property_count
+            Number of node properties.
+        relationship_property_count
+            Number of relationship properties.
+
+
+        Returns
+        -------
+        SessionMemory
+            The estimated memory required for the session.
         """
         if algorithm_categories is None:
             algorithm_categories = []

graphdatascience/tests/unit/test_dedicated_sessions.py

@@ -779,6 +779,13 @@ def test_estimate_size() -> None:
     assert sessions.estimate(1, 1, [AlgorithmCategory.CENTRALITY]) == SessionMemory.m_8GB
 
 
+def test_estimate_str_categories_size() -> None:
+    aura_api = FakeAuraApi(size_estimation=EstimationDetails("1GB", "8GB"))
+    sessions = DedicatedSessions(aura_api)
+
+    assert sessions.estimate(1, 1, ["centrality"]) == SessionMemory.m_8GB
+
+
 def test_estimate_size_exceeds() -> None:
     aura_api = FakeAuraApi(size_estimation=EstimationDetails("16GB", "8GB"))
     sessions = DedicatedSessions(aura_api)