Update generator.py (#206)

BioDataFuse · Jan 31, 2025 · 43123c7 · 43123c7
2 parents 0301cd8 + 5972c2f
commit 43123c7
Showing 1 changed file with 4 additions and 0 deletions.
diff --git a/src/pyBiodatafuse/graph/generator.py b/src/pyBiodatafuse/graph/generator.py
@@ -939,6 +939,8 @@ def save_graph(
     :param disease_compound: the input DataFrame containing disease-compound relationships.
     :param graph_name: the name of the graph.
     :param graph_dir: the directory to save the graph.
+    :returns: a NetworkX MultiDiGraph
+
     """
     graph_path = f"{graph_dir}/{graph_name}"
     os.makedirs(graph_path, exist_ok=True)
@@ -965,3 +967,5 @@ def save_graph(
         pickle.dump(g, f)
     nx.write_gml(g, graph_path_gml)
     logger.warning(f"Graph saved in {graph_path_pickle} and {graph_path_gml}")
+
+    return g