Fix typing issues

Author: s1ck

graphdatascience/error/endpoint_suggester.py

@@ -9,7 +9,7 @@ def generate_suggestive_error_message(requested_endpoint: str, all_endpoints: Li
     MIN_SIMILARITY_FOR_SUGGESTION = 0.9
 
     closest_endpoint = None
-    curr_max_similarity = 0
+    curr_max_similarity = 0.0
     for ep in all_endpoints:
         similarity = textdistance.jaro_winkler(requested_endpoint, ep)
         if similarity >= MIN_SIMILARITY_FOR_SUGGESTION:

graphdatascience/query_runner/arrow_endpoint_version.py

@@ -1,35 +1,35 @@
 from __future__ import annotations
 
 from enum import Enum
+from typing import List
 
 
 class ArrowEndpointVersion(Enum):
     ALPHA = ""
     V1 = "v1/"
 
-    def name(self) -> str:
+    def version(self) -> str:
         return self._name_.lower()
 
     def prefix(self) -> str:
         return self._value_
 
     @staticmethod
-    def from_arrow_info(arrow_info: Series[Any]) -> ArrowEndpointVersion:
-        supported_arrow_versions = arrow_info.get("versions", [])
+    def from_arrow_info(supported_arrow_versions: List[str]) -> ArrowEndpointVersion:
         # Fallback for pre 2.6.0 servers that do not support versions
         if len(supported_arrow_versions) == 0:
             return ArrowEndpointVersion.ALPHA
 
         # If the server supports versioned endpoints, we try v1 first
-        if ArrowEndpointVersion.V1.name() in supported_arrow_versions:
+        if ArrowEndpointVersion.V1.version() in supported_arrow_versions:
             return ArrowEndpointVersion.V1
 
-        if ArrowEndpointVersion.ALPHA.name() in supported_arrow_versions:
+        if ArrowEndpointVersion.ALPHA.version() in supported_arrow_versions:
             return ArrowEndpointVersion.ALPHA
 
         raise UnsupportedArrowEndpointVersion(supported_arrow_versions)
 
 
 class UnsupportedArrowEndpointVersion(Exception):
-    def __init__(self, server_version):
+    def __init__(self, server_version: List[str]) -> None:
         super().__init__(self, f"Unsupported Arrow endpoint versions: {server_version}")

graphdatascience/query_runner/arrow_graph_constructor.py

@@ -126,13 +126,13 @@ def _send_dfs(self, dfs: List[DataFrame], entity_type: str) -> None:
     def _versioned_action_type(self, action_type: str) -> str:
         return self._arrow_endpoint_version.prefix() + action_type
 
-    def _versioned_flight_desriptor(self, flight_descriptor: dict) -> dict:
+    def _versioned_flight_desriptor(self, flight_descriptor: Dict[str, Any]) -> Dict[str, Any]:
         return (
             flight_descriptor
             if self._arrow_endpoint_version == ArrowEndpointVersion.ALPHA
             else {
                 "name": "PUT_MESSAGE",
-                "version": ArrowEndpointVersion.V1.name(),
+                "version": ArrowEndpointVersion.V1.version(),
                 "body": flight_descriptor,
             }
         )

graphdatascience/query_runner/arrow_query_runner.py

@@ -7,7 +7,7 @@
 from typing import Any, Dict, List, Optional, Tuple
 
 import pyarrow.flight as flight
-from pandas import DataFrame, Series
+from pandas import DataFrame
 from pyarrow import ChunkedArray, Table, chunked_array
 from pyarrow.flight import ClientMiddleware, ClientMiddlewareFactory
 from pyarrow.types import is_dictionary  # type: ignore
@@ -31,11 +31,13 @@ def create(
         encrypted: bool = False,
         disable_server_verification: bool = False,
         tls_root_certs: Optional[bytes] = None,
-    ) -> "QueryRunner":
-        arrow_info = ArrowQueryRunner._get_arrow_debug_info(fallback_query_runner)
+    ) -> QueryRunner:
+        arrow_info = (
+            fallback_query_runner.call_procedure(endpoint="gds.debug.arrow", custom_error=False).squeeze().to_dict()
+        )
         server_version = fallback_query_runner.server_version()
-        listen_address: str = arrow_info.get("advertisedListenAddress", arrow_info["listenAddress"])  # type: ignore
-        arrow_endpoint_version = ArrowEndpointVersion.from_arrow_info(arrow_info)
+        listen_address: str = arrow_info.get("advertisedListenAddress", arrow_info["listenAddress"])
+        arrow_endpoint_version = ArrowEndpointVersion.from_arrow_info(arrow_info.get("versions", []))
 
         if arrow_info["running"]:
             return ArrowQueryRunner(
@@ -51,10 +53,6 @@ def create(
         else:
             return fallback_query_runner
 
-    @staticmethod
-    def _get_arrow_debug_info(query_runner) -> Series[Any]:
-        return query_runner.call_procedure(endpoint="gds.debug.arrow", custom_error=False).squeeze()
-
     def __init__(
         self,
         uri: str,
@@ -280,7 +278,7 @@ def _run_arrow_property_get(self, graph_name: str, procedure_name: str, configur
         if self._arrow_endpoint_version == ArrowEndpointVersion.V1:
             payload = {
                 "name": "GET_MESSAGE",
-                "version": ArrowEndpointVersion.V1.name(),
+                "version": ArrowEndpointVersion.V1.version(),
                 "body": payload,
             }
 

graphdatascience/tests/unit/test_arrow_endpoint_version.py

@@ -1,6 +1,7 @@
+from typing import List, Tuple
+
 import pytest
 
-from graphdatascience.graph_data_science import ArrowQueryRunner
 from graphdatascience.query_runner.arrow_endpoint_version import (
     ArrowEndpointVersion,
     UnsupportedArrowEndpointVersion,
@@ -19,23 +20,19 @@
         (ArrowEndpointVersion.ALPHA, ["alpha", "v2"]),
     ],
 )
-def test_from_arrow_info_multiple_versions(arrow_versions):
-    arrow_info = dict()
-    arrow_info["versions"] = arrow_versions[1]
-    actual = ArrowEndpointVersion.from_arrow_info(arrow_info)
+def test_from_arrow_info_multiple_versions(arrow_versions: Tuple[ArrowEndpointVersion, List[str]]) -> None:
+    actual = ArrowEndpointVersion.from_arrow_info(arrow_versions[1])
     assert actual == arrow_versions[0]
 
 
 @pytest.mark.parametrize("arrow_version", ["v2", "unsupported"])
-def test_from_arrow_info_fails(arrow_version):
-    arrow_info = dict()
-    arrow_info["versions"] = [arrow_version]
+def test_from_arrow_info_fails(arrow_version: str) -> None:
     with pytest.raises(UnsupportedArrowEndpointVersion) as e:
-        ArrowEndpointVersion.from_arrow_info(arrow_info)
+        ArrowEndpointVersion.from_arrow_info([arrow_version])
     assert "Unsupported" in str(e.value)
     assert arrow_version in str(e.value)
 
 
-def test_prefix():
+def test_prefix() -> None:
     assert ArrowEndpointVersion.ALPHA.prefix() == ""
     assert ArrowEndpointVersion.V1.prefix() == "v1/"