Fix codestyle

Author: DarthMax

graphdatascience/arrow_client/v2/job_client.py

@@ -45,7 +45,7 @@ def get_summary(client: AuthenticatedArrowClient, job_id: str) -> dict[str, Any]
         return deserialize_single(res)
 
     @staticmethod
-    def stream_results(client: AuthenticatedArrowClient,graph_name: str, job_id: str) -> DataFrame:
+    def stream_results(client: AuthenticatedArrowClient, graph_name: str, job_id: str) -> DataFrame:
         payload = {
             "graphName": graph_name,
             "jobId": job_id,
@@ -54,15 +54,10 @@ def stream_results(client: AuthenticatedArrowClient,graph_name: str, job_id: str
         res = client.do_action_with_retry("v2/results.stream", json.dumps(payload).encode("utf-8"))
         export_job_id = JobIdConfig(**deserialize_single(res)).job_id
 
-        payload = {
-            "version": "v2",
-            "name": export_job_id,
-            "body": {}
-        }
+        stream_payload = {"version": "v2", "name": export_job_id, "body": {}}
 
-        ticket = Ticket(json.dumps(payload).encode("utf-8"))
+        ticket = Ticket(json.dumps(stream_payload).encode("utf-8"))
 
         get = client.get_stream(ticket)
         arrow_table = get.read_all()
         return arrow_table.to_pandas(types_mapper=ArrowDtype)  # type: ignore
-

graphdatascience/procedure_surface/arrow/arrow_wcc_endpoints.py

@@ -174,8 +174,7 @@ def write(
 
         return WccWriteResult(**computation_result)
 
-    def estimate(self, graph_name: Optional[str] = None,
-                 projection_config: Optional[dict[str, Any]] = None) -> EstimationResult:
+    def estimate(
+        self, graph_name: Optional[str] = None, projection_config: Optional[dict[str, Any]] = None
+    ) -> EstimationResult:
         pass
-
-

graphdatascience/tests/integrationV2/procedure_surface/arrow/conftest.py

@@ -8,13 +8,12 @@
 from testcontainers.neo4j import Neo4jContainer
 
 from graphdatascience.arrow_client.arrow_authentication import UsernamePasswordAuthentication
+from graphdatascience.arrow_client.arrow_info import ArrowInfo
 from graphdatascience.arrow_client.authenticated_flight_client import AuthenticatedArrowClient
-from graphdatascience.arrow_client.v2.write_back_client import WriteBackClient
-from graphdatascience.query_runner.arrow_info import ArrowInfo
 
 
 @pytest.fixture(scope="session")
-def password_file():
+def password_file() -> Generator[str, None, None]:
     """Create a temporary file and return its path."""
     temp_dir = tempfile.mkdtemp()
     temp_file_path = os.path.join(temp_dir, "password")
@@ -84,12 +83,3 @@ def arrow_client(session_container: DockerContainer) -> AuthenticatedArrowClient
         auth=UsernamePasswordAuthentication("neo4j", "password"),
         encrypted=False,
     )
-
-
-@pytest.fixture
-def write_back_client(neo4j_session_container: DockerContainer) -> WriteBackClient:
-    """Create a write-back client for the session container."""
-    host = neo4j_session_container.get_container_host_ip()
-    port = neo4j_session_container.get_exposed_port(8491)
-
-    return WriteBackClient(host=host, port=port, username="neo4j", password="test_password")

graphdatascience/tests/integrationV2/procedure_surface/arrow/test_wcc_arrow_endpoints.py

@@ -1,10 +1,10 @@
 import json
+from typing import Generator
 
 import pytest
 
 from graphdatascience import Graph
 from graphdatascience.arrow_client.authenticated_flight_client import AuthenticatedArrowClient
-from graphdatascience.arrow_client.v2.data_mapper_utils import deserialize_single
 from graphdatascience.procedure_surface.arrow.arrow_wcc_endpoints import WccArrowEndpoints
 
 
@@ -17,36 +17,36 @@ def name(self) -> str:
 
 
 @pytest.fixture
-def sample_graph(arrow_client: AuthenticatedArrowClient):
+def sample_graph(arrow_client: AuthenticatedArrowClient) -> Generator[Graph, None, None]:
     gdl = """
     (a: Node)
     (b: Node)
     (c: Node)
     (a)-[:REL]->(c)
     """
 
-    arrow_client.do_action( "v2/graph.fromGDL", json.dumps({"graphName": "g", "gdlGraph": gdl}).encode("utf-8"))
+    arrow_client.do_action("v2/graph.fromGDL", json.dumps({"graphName": "g", "gdlGraph": gdl}).encode("utf-8"))
     yield MockGraph("g")
-    arrow_client.do_action( "v2/graph.drop", json.dumps({"graphName": "g"}).encode("utf-8"))
+    arrow_client.do_action("v2/graph.drop", json.dumps({"graphName": "g"}).encode("utf-8"))
+
 
 @pytest.fixture
-def wcc_endpoints(arrow_client: AuthenticatedArrowClient):
+def wcc_endpoints(arrow_client: AuthenticatedArrowClient) -> Generator[WccArrowEndpoints, None, None]:
     yield WccArrowEndpoints(arrow_client)
 
 
-def test_wcc_stats(wcc_endpoints: WccArrowEndpoints, sample_graph: Graph):
+def test_wcc_stats(wcc_endpoints: WccArrowEndpoints, sample_graph: Graph) -> None:
     """Test WCC stats operation."""
-    result = wcc_endpoints.stats(
-        G=sample_graph
-    )
+    result = wcc_endpoints.stats(G=sample_graph)
 
     assert result.component_count == 2
     assert result.compute_millis > 0
     assert result.pre_processing_millis > 0
     assert result.post_processing_millis > 0
     assert "p10" in result.component_distribution
 
-def test_wcc_stream(wcc_endpoints: WccArrowEndpoints, sample_graph: Graph):
+
+def test_wcc_stream(wcc_endpoints: WccArrowEndpoints, sample_graph: Graph) -> None:
     """Test WCC stream operation."""
     result_df = wcc_endpoints.stream(
         G=sample_graph,
@@ -56,7 +56,8 @@ def test_wcc_stream(wcc_endpoints: WccArrowEndpoints, sample_graph: Graph):
     assert "componentId" in result_df.columns
     assert len(result_df.columns) == 2
 
-def test_wcc_mutate(wcc_endpoints: WccArrowEndpoints, sample_graph: Graph):
+
+def test_wcc_mutate(wcc_endpoints: WccArrowEndpoints, sample_graph: Graph) -> None:
     """Test WCC mutate operation."""
     result = wcc_endpoints.mutate(
         G=sample_graph,

graphdatascience/tests/integrationV2/procedure_surface/cypher/conftest.py

@@ -6,6 +6,7 @@
 from testcontainers.core.waiting_utils import wait_for_logs
 from testcontainers.neo4j import Neo4jContainer
 
+from graphdatascience import QueryRunner
 from graphdatascience.query_runner.neo4j_query_runner import Neo4jQueryRunner
 
 
@@ -27,7 +28,7 @@ def neo4j_database_container() -> Generator[Neo4jContainer, None, None]:
 
 
 @pytest.fixture
-def query_runner(neo4j_database_container: DockerContainer):
+def query_runner(neo4j_database_container: DockerContainer) -> Generator[QueryRunner, None, None]:
     yield Neo4jQueryRunner.create_for_db(
         f"bolt://localhost:{neo4j_database_container.get_exposed_port(7687)}",
         ("neo4j", "password"),

graphdatascience/tests/integrationV2/procedure_surface/cypher/test_wcc_cypher_endpoints.py

@@ -1,13 +1,14 @@
+from typing import Generator
+
 import pytest
 
 from graphdatascience import Graph, QueryRunner
 from graphdatascience.procedure_surface.arrow.arrow_wcc_endpoints import WccArrowEndpoints
 from graphdatascience.procedure_surface.cypher.wcc_cypher_endpoints import WccCypherEndpoints
-from graphdatascience.tests.integrationV2.procedure_surface.cypher.conftest import query_runner
 
 
 @pytest.fixture
-def sample_graph(query_runner: QueryRunner):
+def sample_graph(query_runner: QueryRunner) -> Generator[Graph, None, None]:
     create_statement = """
     CREATE
     (a: Node),
@@ -30,24 +31,24 @@ def sample_graph(query_runner: QueryRunner):
     query_runner.run_cypher("CALL gds.graph.drop('g')")
     query_runner.run_cypher("MATCH (n) DETACH DELETE n")
 
+
 @pytest.fixture
-def wcc_endpoints(query_runner: QueryRunner):
+def wcc_endpoints(query_runner: QueryRunner) -> Generator[WccCypherEndpoints, None, None]:
     yield WccCypherEndpoints(query_runner)
 
 
-def test_wcc_stats(wcc_endpoints: WccArrowEndpoints, sample_graph: Graph):
+def test_wcc_stats(wcc_endpoints: WccArrowEndpoints, sample_graph: Graph) -> None:
     """Test WCC stats operation."""
-    result = wcc_endpoints.stats(
-        G=sample_graph
-    )
+    result = wcc_endpoints.stats(G=sample_graph)
 
     assert result.component_count == 2
     assert result.compute_millis > 0
     assert result.pre_processing_millis > 0
     assert result.post_processing_millis > 0
     assert "p10" in result.component_distribution
 
-def test_wcc_stream(wcc_endpoints: WccArrowEndpoints, sample_graph: Graph):
+
+def test_wcc_stream(wcc_endpoints: WccArrowEndpoints, sample_graph: Graph) -> None:
     """Test WCC stream operation."""
     result_df = wcc_endpoints.stream(
         G=sample_graph,
@@ -57,7 +58,8 @@ def test_wcc_stream(wcc_endpoints: WccArrowEndpoints, sample_graph: Graph):
     assert "componentId" in result_df.columns
     assert len(result_df.columns) == 2
 
-def test_wcc_mutate(wcc_endpoints: WccArrowEndpoints, sample_graph: Graph):
+
+def test_wcc_mutate(wcc_endpoints: WccArrowEndpoints, sample_graph: Graph) -> None:
     """Test WCC mutate operation."""
     result = wcc_endpoints.mutate(
         G=sample_graph,