Fix type issues with ndarray.tolist()

Author: FlorentinD

graphdatascience/graph/ogb_loader.py

@@ -108,12 +108,12 @@ def _parse_homogeneous(self, dataset: HomogeneousOGBNDataset) -> tuple[list[pd.D
             "nodeId": list(range(node_count)),
         }
         if "node_feat" in graph and graph["node_feat"] is not None:
-            node_dict["features"] = graph["node_feat"].tolist()
+            node_dict["features"] = graph["node_feat"].tolist()  # type: ignore
 
         if len(dataset.labels[0]) == 1:
             node_dict["classLabel"] = [cl[0] for cl in dataset.labels]
         else:
-            node_dict["classLabel"] = dataset.labels.tolist()
+            node_dict["classLabel"] = dataset.labels.tolist()  # type: ignore
 
         split = dataset.get_idx_split()
         node_labels = ["Train" for _ in range(node_count)]
@@ -170,13 +170,13 @@ def _parse_heterogeneous(self, dataset: HeterogeneousOGBNDataset) -> tuple[list[
             }
 
             if node_label in node_features:
-                node_dict["features"] = node_features[node_label].tolist()
+                node_dict["features"] = node_features[node_label].tolist()  # type: ignore
 
             if node_label in class_labels:
                 if len(class_labels[node_label]) == 1:
                     node_dict["classLabel"] = [cl[0] for cl in class_labels[node_label]]
                 else:
-                    node_dict["classLabel"] = class_labels[node_label].tolist()
+                    node_dict["classLabel"] = class_labels[node_label].tolist()  # type: ignore
 
             node_id_offsets[node_label] = current_offset
             current_offset += node_count
@@ -243,7 +243,7 @@ def _parse_homogeneous(self, dataset: HomogeneousOGBLDataset) -> tuple[list[pd.D
             "labels": "N",
         }
         if "node_feat" in graph and graph["node_feat"] is not None:
-            node_dict["features"] = graph["node_feat"].tolist()
+            node_dict["features"] = graph["node_feat"].tolist()  # type: ignore
         nodes = pd.DataFrame(node_dict)
 
         self._logger.info("Preparing relationship data for transfer to server...")
@@ -283,7 +283,7 @@ def _parse_heterogeneous(self, dataset: HeterogeneousOGBLDataset) -> tuple[list[
             }
 
             if node_label in node_features:
-                node_dict["features"] = node_features[node_label].tolist()
+                node_dict["features"] = node_features[node_label].tolist()  # type: ignore
 
             node_id_offsets[node_label] = current_offset
             current_offset += node_count

graphdatascience/query_runner/cypher_graph_constructor.py

@@ -386,7 +386,8 @@ def run(self, node_df: DataFrame, relationship_df: DataFrame) -> None:
             )
 
         def _node_query(self, node_df: DataFrame) -> tuple[str, list[list[Any]]]:
-            node_list = node_df.values.tolist()
+            # ignore type as tolist return depends on number of dimensions)
+            node_list: list[list[Any]] = node_df.values.tolist()  # type: ignore
             node_columns = list(node_df.columns)
             node_id_index = node_columns.index("nodeId")
 
@@ -411,7 +412,7 @@ def _node_query(self, node_df: DataFrame) -> tuple[str, list[list[Any]]]:
             return f"UNWIND $nodes as node RETURN node[{node_id_index}] as id{label_query}{property_query}", node_list
 
         def _relationship_query(self, rel_df: DataFrame) -> tuple[str, list[list[Any]]]:
-            rel_list = rel_df.values.tolist()
+            rel_list: list[list[Any]] = rel_df.values.tolist()  # type: ignore
             rel_columns = list(rel_df.columns)
             source_id_index = rel_columns.index("sourceNodeId")
             target_id_index = rel_columns.index("targetNodeId")