handle None location case and remove print statements

funkelab · Mar 13, 2024 · 8db2a23 · 8db2a23
1 parent ea69950
commit 8db2a23
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 9 deletions.
diff --git a/funlib/persistence/graphs/pgsql_graph_database.py b/funlib/persistence/graphs/pgsql_graph_database.py
@@ -137,7 +137,6 @@ def _read_metadata(self) -> Optional[dict[str, Any]]:
         return None
 
     def _select_query(self, query) -> Iterable[Any]:
-        print(query)
         self.__exec(query)
         return self.cur
 

diff --git a/funlib/persistence/graphs/sql_graph_database.py b/funlib/persistence/graphs/sql_graph_database.py
@@ -345,11 +345,6 @@ def read_nodes(
             for values in self._select_query(select_statement)
         ]
 
-        for values in self._select_query(select_statement):
-            print(values)
-        print(self.node_attrs.keys())
-        print(nodes)
-
         return nodes
 
     def num_nodes(self, roi: Roi) -> int:
@@ -445,7 +440,7 @@ def write_edges(
                 u, v = min(u, v), max(u, v)
             pos_u = self.__get_node_pos(nodes[u])
 
-            if not roi.contains(pos_u):
+            if pos_u is None or not roi.contains(pos_u):
                 logger.debug(
                     (
                         f"Skipping edge with {self.endpoint_names[0]} {{}}, {self.endpoint_names[1]} {{}},"
@@ -660,8 +655,11 @@ def __remove_keys(self, dictionary, keys):
 
         return {k: v for k, v in dictionary.items() if k not in keys}
 
-    def __get_node_pos(self, n: dict[str, Any]) -> Coordinate:
-        return Coordinate(n[self.position_attribute])
+    def __get_node_pos(self, n: dict[str, Any]) -> Optional[Coordinate]:
+        try:
+            return Coordinate(n[self.position_attribute])
+        except KeyError:
+            return None
 
     def __convert_to_sql(self, x: Any) -> str:
         if isinstance(x, str):