Add type hints

neo4j · Dec 21, 2022 · 32cce59 · 32cce59
1 parent d72ea93
commit 32cce59
Showing 1 changed file with 4 additions and 4 deletions.
diff --git 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)