Run black

Author: adsharma

src_py/async_connection.py

@@ -168,13 +168,17 @@ async def execute(
             conn, conn_index = self.__get_connection_with_least_queries()
 
         try:
-            return await loop.run_in_executor(self.executor, conn.execute, query, parameters)
+            return await loop.run_in_executor(
+                self.executor, conn.execute, query, parameters
+            )
         except asyncio.CancelledError:
             conn.interrupt()
         finally:
             self.__decrement_connection_counter(conn_index)
 
-    async def _prepare(self, query: str, parameters: dict[str, Any] | None = None) -> PreparedStatement:
+    async def _prepare(
+        self, query: str, parameters: dict[str, Any] | None = None
+    ) -> PreparedStatement:
         """
         The only parameters supported during prepare are dataframes.
         Any remaining parameters will be ignored and should be passed to execute().
@@ -183,12 +187,16 @@ async def _prepare(self, query: str, parameters: dict[str, Any] | None = None) -
         conn, conn_index = self.__get_connection_with_least_queries()
 
         try:
-            preparedStatement = await loop.run_in_executor(self.executor, conn.prepare, query, parameters)
+            preparedStatement = await loop.run_in_executor(
+                self.executor, conn.prepare, query, parameters
+            )
             return preparedStatement
         finally:
             self.__decrement_connection_counter(conn_index)
 
-    async def prepare(self, query: str, parameters: dict[str, Any] | None = None) -> PreparedStatement:
+    async def prepare(
+        self, query: str, parameters: dict[str, Any] | None = None
+    ) -> PreparedStatement:
         """
         Create a prepared statement for a query asynchronously.
 

src_py/connection.py

@@ -1,7 +1,8 @@
 from __future__ import annotations
 
 import warnings
-from typing import TYPE_CHECKING, Any, Callable
+from collections.abc import Callable
+from typing import TYPE_CHECKING, Any
 
 from . import _lbug
 from .prepared_statement import PreparedStatement
@@ -130,8 +131,12 @@ def execute(
         if len(parameters) == 0 and isinstance(query, str):
             query_result_internal = self._connection.query(query)
         else:
-            prepared_statement = self._prepare(query, parameters) if isinstance(query, str) else query
-            query_result_internal = self._connection.execute(prepared_statement._prepared_statement, parameters)
+            prepared_statement = (
+                self._prepare(query, parameters) if isinstance(query, str) else query
+            )
+            query_result_internal = self._connection.execute(
+                prepared_statement._prepared_statement, parameters
+            )
         if not query_result_internal.isSuccess():
             raise RuntimeError(query_result_internal.getErrorMessage())
         current_query_result = QueryResult(self, query_result_internal)
@@ -234,7 +239,9 @@ def _get_rel_table_names(self) -> list[dict[str, Any]]:
             row = tables_result.get_next()
             if row[2] == "REL":
                 name = row[1]
-                connections_result = self.execute(f"CALL show_connection({name!r}) RETURN *;")
+                connections_result = self.execute(
+                    f"CALL show_connection({name!r}) RETURN *;"
+                )
                 src_dst_row = connections_result.get_next()
                 src_node = src_dst_row[0]
                 dst_node = src_dst_row[1]
@@ -345,7 +352,9 @@ def create_arrow_table(
 
         """
         self.init_connection()
-        query_result_internal = self._connection.create_arrow_table(table_name, dataframe)
+        query_result_internal = self._connection.create_arrow_table(
+            table_name, dataframe
+        )
         if not query_result_internal.isSuccess():
             raise RuntimeError(query_result_internal.getErrorMessage())
         return QueryResult(self, query_result_internal)

src_py/database.py

@@ -253,19 +253,29 @@ def _scan_node_table(
 
         if prop_type == Type.INT64.value:
             result = np.empty(len(indices) * dim, dtype=np.int64)
-            self._database.scan_node_table_as_int64(table_name, prop_name, indices_cast, result, num_threads)
+            self._database.scan_node_table_as_int64(
+                table_name, prop_name, indices_cast, result, num_threads
+            )
         elif prop_type == Type.INT32.value:
             result = np.empty(len(indices) * dim, dtype=np.int32)
-            self._database.scan_node_table_as_int32(table_name, prop_name, indices_cast, result, num_threads)
+            self._database.scan_node_table_as_int32(
+                table_name, prop_name, indices_cast, result, num_threads
+            )
         elif prop_type == Type.INT16.value:
             result = np.empty(len(indices) * dim, dtype=np.int16)
-            self._database.scan_node_table_as_int16(table_name, prop_name, indices_cast, result, num_threads)
+            self._database.scan_node_table_as_int16(
+                table_name, prop_name, indices_cast, result, num_threads
+            )
         elif prop_type == Type.DOUBLE.value:
             result = np.empty(len(indices) * dim, dtype=np.float64)
-            self._database.scan_node_table_as_double(table_name, prop_name, indices_cast, result, num_threads)
+            self._database.scan_node_table_as_double(
+                table_name, prop_name, indices_cast, result, num_threads
+            )
         elif prop_type == Type.FLOAT.value:
             result = np.empty(len(indices) * dim, dtype=np.float32)
-            self._database.scan_node_table_as_float(table_name, prop_name, indices_cast, result, num_threads)
+            self._database.scan_node_table_as_float(
+                table_name, prop_name, indices_cast, result, num_threads
+            )
 
         if result is not None:
             return result

src_py/prepared_statement.py

@@ -12,7 +12,12 @@ class PreparedStatement:
     same query for repeated execution.
     """
 
-    def __init__(self, connection: Connection, query: str, parameters: dict[str, Any] | None = None):
+    def __init__(
+        self,
+        connection: Connection,
+        query: str,
+        parameters: dict[str, Any] | None = None,
+    ):
         """
         Parameters
         ----------

src_py/query_result.py

@@ -2,11 +2,10 @@
 
 from typing import TYPE_CHECKING
 
+from .constants import DST, ID, LABEL, NODES, RELS, SRC
 from .torch_geometric_result_converter import TorchGeometricResultConverter
 from .types import Type
 
-from .constants import ID, LABEL, SRC, DST, NODES, RELS
-
 if TYPE_CHECKING:
     import sys
     from collections.abc import Iterator
@@ -191,7 +190,9 @@ def get_as_pl(self) -> pl.DataFrame:
             data=self.get_as_arrow(chunk_size=-1, fallbackExtensionTypes=True),
         )
 
-    def get_as_arrow(self, chunk_size: int | None = None, *, fallbackExtensionTypes: bool = False) -> pa.Table:
+    def get_as_arrow(
+        self, chunk_size: int | None = None, *, fallbackExtensionTypes: bool = False
+    ) -> pa.Table:
         """
         Get the query result as a PyArrow Table.
 
@@ -314,7 +315,9 @@ def get_as_networkx(
         table_to_label_dict = {}
         table_primary_key_dict = {}
 
-        def encode_node_id(node: dict[str, Any], table_primary_key_dict: dict[str, Any]) -> str:
+        def encode_node_id(
+            node: dict[str, Any], table_primary_key_dict: dict[str, Any]
+        ) -> str:
             node_label = node[LABEL]
             return f"{node_label}_{node[table_primary_key_dict[node_label]]!s}"
 

src_py/torch_geometric_feature_store.py

@@ -68,20 +68,29 @@ def __get_tensor_by_scan(self, attr: TensorAttr) -> FeatureTensorType | None:
             if indices.step is None or indices.step == 1:
                 indices = np.arange(indices.start, indices.stop, dtype=np.uint64)
             else:
-                indices = np.arange(indices.start, indices.stop, indices.step, dtype=np.uint64)
+                indices = np.arange(
+                    indices.start, indices.stop, indices.step, dtype=np.uint64
+                )
         elif isinstance(indices, int):
             indices = np.array([indices])
 
         if table_name not in self.node_properties_cache:
-            self.node_properties_cache[table_name] = self.connection._get_node_property_names(table_name)
+            self.node_properties_cache[table_name] = (
+                self.connection._get_node_property_names(table_name)
+            )
         attr_info = self.node_properties_cache[table_name][attr_name]
 
         flat_dim = 1
         if attr_info["dimension"] > 0:
             for i in range(attr_info["dimension"]):
                 flat_dim *= attr_info["shape"][i]
         scan_result = self.connection.database._scan_node_table(
-            table_name, attr_name, attr_info["type"], flat_dim, indices, self.num_threads
+            table_name,
+            attr_name,
+            attr_info["type"],
+            flat_dim,
+            indices,
+            self.num_threads,
         )
 
         if attr_info["dimension"] > 0 and "shape" in attr_info:
@@ -151,11 +160,16 @@ def _get_tensor_size(self, attr: TensorAttr) -> tuple[Any, ...]:
             return (length,) + attr_info["shape"]
 
     def __get_node_property(self, table_name: str, attr_name: str) -> dict[str, Any]:
-        if table_name in self.node_properties_cache and attr_name in self.node_properties_cache[table_name]:
+        if (
+            table_name in self.node_properties_cache
+            and attr_name in self.node_properties_cache[table_name]
+        ):
             return self.node_properties_cache[table_name][attr_name]
         self.__get_connection()
         if table_name not in self.node_properties_cache:
-            self.node_properties_cache[table_name] = self.connection._get_node_property_names(table_name)
+            self.node_properties_cache[table_name] = (
+                self.connection._get_node_property_names(table_name)
+            )
         if attr_name not in self.node_properties_cache[table_name]:
             msg = f"Attribute {attr_name} not found in group {table_name}"
             raise ValueError(msg)
@@ -168,7 +182,9 @@ def get_all_tensor_attrs(self) -> list[TensorAttr]:
         self.__get_connection()
         for table_name in self.connection._get_node_table_names():
             if table_name not in self.node_properties_cache:
-                self.node_properties_cache[table_name] = self.connection._get_node_property_names(table_name)
+                self.node_properties_cache[table_name] = (
+                    self.connection._get_node_property_names(table_name)
+                )
             for attr_name in self.node_properties_cache[table_name]:
                 if self.node_properties_cache[table_name][attr_name]["type"] in [
                     Type.INT64.value,

src_py/torch_geometric_graph_store.py

@@ -20,7 +20,7 @@
     if sys.version_info >= (3, 10):
         from typing import TypeAlias
     else:
-        from typing_extensions import TypeAlias
+        from typing import TypeAlias
 
 StoreKeyType: TypeAlias = tuple[tuple[str], Any, bool]
 
@@ -61,7 +61,9 @@ def _put_edge_index(self, edge_index: EdgeTensorType, edge_attr: EdgeAttr) -> No
             self.store[key].materialized = True
             self.store[key].size = edge_attr.size
         else:
-            self.store[key] = Rel(key[0], key[1], key[2], edge_attr.size, True, edge_index)
+            self.store[key] = Rel(
+                key[0], key[1], key[2], edge_attr.size, True, edge_index
+            )
 
     def _get_edge_index(self, edge_attr: EdgeAttr) -> EdgeTensorType | None:
         if edge_attr.layout.value == EdgeLayout.COO.value:  # noqa: SIM102
@@ -91,7 +93,10 @@ def _remove_edge_index(self, edge_attr: EdgeAttr) -> None:
 
     def get_all_edge_attrs(self) -> list[EdgeAttr]:
         """Return all EdgeAttr from the store values."""
-        return [EdgeAttr(rel.edge_type, rel.layout, rel.is_sorted, rel.size) for rel in self.store.values()]
+        return [
+            EdgeAttr(rel.edge_type, rel.layout, rel.is_sorted, rel.size)
+            for rel in self.store.values()
+        ]
 
     def __get_edge_coo_from_database(self, key: StoreKeyType) -> None:
         if not self.connection: