From 32cce59beb53752d2dec6ac01fd412daf6e6a65e Mon Sep 17 00:00:00 2001 From: Mats Rydberg Date: Wed, 21 Dec 2022 13:11:32 +0100 Subject: [PATCH] Add type hints --- graphdatascience/query_runner/cypher_graph_constructor.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/graphdatascience/query_runner/cypher_graph_constructor.py b/graphdatascience/query_runner/cypher_graph_constructor.py index a217b97b2..ecc6545f4 100644 --- a/graphdatascience/query_runner/cypher_graph_constructor.py +++ b/graphdatascience/query_runner/cypher_graph_constructor.py @@ -222,8 +222,8 @@ def adjust_node_df(self, node_dfs: List[DataFrame], schema: GraphColumnSchema) - for i, df in enumerate(node_dfs): node_dict: Dict[str, Any] = { "sourceNodeId": df["nodeId"], - f"targetNodeId{self._BIT_COL_SUFFIX}": False, "targetNodeId": -1, + f"targetNodeId{self._BIT_COL_SUFFIX}": False, } if CypherAggregationApi.RELATIONSHIP_TYPE in schema.all_rels.all: @@ -237,7 +237,7 @@ def adjust_node_df(self, node_dfs: List[DataFrame], schema: GraphColumnSchema) - node_dict[CypherAggregationApi.SOURCE_NODE_LABEL + self._BIT_COL_SUFFIX] = False node_dict[CypherAggregationApi.SOURCE_NODE_LABEL] = "" - def collect_to_dict(row): # type: ignore + def collect_to_dict(row: Dict[str, Any]) -> Dict[str, Any]: return {column: row[column] for column in schema.nodes_per_df[i].properties} node_dict_df = DataFrame(node_dict) @@ -254,7 +254,7 @@ def adjust_rel_df(self, rel_dfs: List[DataFrame], schema: GraphColumnSchema) -> adjusted_dfs = [] for i, df in enumerate(rel_dfs): - rel_dict = { + rel_dict: Dict[str, Any] = { "sourceNodeId": df["sourceNodeId"], "targetNodeId": df["targetNodeId"], f"targetNodeId{self._BIT_COL_SUFFIX}": True, @@ -271,7 +271,7 @@ def adjust_rel_df(self, rel_dfs: List[DataFrame], schema: GraphColumnSchema) -> rel_dict[CypherAggregationApi.SOURCE_NODE_LABEL] = None rel_dict[CypherAggregationApi.SOURCE_NODE_LABEL + self._BIT_COL_SUFFIX] = False - def collect_to_dict(row): # type: ignore + def collect_to_dict(row: Dict[str, Any]) -> Dict[str, Any]: return {column: row[column] for column in schema.rels_per_df[i].properties} rel_dict_df = DataFrame(rel_dict)