graphdatascience 1.17__tar.gz → 1.18__tar.gz

graphdatascience-1.18/MANIFEST.in

@@ -0,0 +1,3 @@
+prune graphdatascience/resources/cora/serialize_cora.py
+prune graphdatascience/resources/imdb/serialize_imdb.py
+prune graphdatascience/resources/lastfm/serialize_lastfm.py

{graphdatascience-1.17/graphdatascience.egg-info → graphdatascience-1.18}/PKG-INFO

@@ -1,22 +1,19 @@
 Metadata-Version: 2.4
 Name: graphdatascience
-Version: 1.17
+Version: 1.18
 Summary: A Python client for the Neo4j Graph Data Science (GDS) library
-Home-page: https://neo4j.com/product/graph-data-science/
-Author: Neo4j
-Author-email: team-gds@neo4j.org
-License: Apache License 2.0
+Author-email: Neo4j <team-gds@neo4j.org>
+License-Expression: Apache-2.0
+Project-URL: Homepage, https://neo4j.com/product/graph-data-science/
 Project-URL: Documentation, https://neo4j.com/docs/graph-data-science-client/current/
-Project-URL: Source, https://github.com/neo4j/graph-data-science-client
-Project-URL: Bug Tracker, https://github.com/neo4j/graph-data-science-client/issues
+Project-URL: Repository, https://github.com/neo4j/graph-data-science-client
+Project-URL: BugTracker, https://github.com/neo4j/graph-data-science-client/issues
 Classifier: Development Status :: 5 - Production/Stable
 Classifier: Intended Audience :: Developers
 Classifier: Intended Audience :: Science/Research
-Classifier: License :: OSI Approved :: Apache Software License
 Classifier: Operating System :: OS Independent
 Classifier: Programming Language :: Python :: 3
 Classifier: Programming Language :: Python :: 3 :: Only
-Classifier: Programming Language :: Python :: 3.9
 Classifier: Programming Language :: Python :: 3.10
 Classifier: Programming Language :: Python :: 3.11
 Classifier: Programming Language :: Python :: 3.12
@@ -25,39 +22,27 @@ Classifier: Topic :: Database
 Classifier: Topic :: Scientific/Engineering
 Classifier: Topic :: Software Development
 Classifier: Typing :: Typed
-Requires-Python: >=3.9
+Requires-Python: >=3.10
 Description-Content-Type: text/markdown
 License-File: LICENSE
 Requires-Dist: multimethod<3.0,>=1.0
-Requires-Dist: neo4j<6.0,>=4.4.12
-Requires-Dist: numpy<2.3
+Requires-Dist: neo4j<7.0,>=4.4.12
+Requires-Dist: numpy<2.4
 Requires-Dist: pandas<3.0,>=1.0
-Requires-Dist: pyarrow<21.0,>=17.0
+Requires-Dist: pyarrow<23.0,>=17.0
 Requires-Dist: textdistance<5.0,>=4.0
 Requires-Dist: tqdm<5.0,>=4.0
 Requires-Dist: typing-extensions<5.0,>=4.0
 Requires-Dist: requests
 Requires-Dist: tenacity>=9.0
 Requires-Dist: pydantic>=2.11
-Provides-Extra: ogb
-Requires-Dist: ogb<2.0,>=1.0; extra == "ogb"
 Provides-Extra: networkx
 Requires-Dist: networkx<4.0,>=2.0; extra == "networkx"
+Provides-Extra: ogb
+Requires-Dist: ogb<2.0,>=1.0; extra == "ogb"
 Provides-Extra: rust-ext
-Requires-Dist: neo4j-rust-ext<6.0,>=4.4.12; extra == "rust-ext"
-Dynamic: author
-Dynamic: author-email
-Dynamic: classifier
-Dynamic: description
-Dynamic: description-content-type
-Dynamic: home-page
-Dynamic: license
+Requires-Dist: neo4j-rust-ext<7.0,>=4.4.12; extra == "rust-ext"
 Dynamic: license-file
-Dynamic: project-url
-Dynamic: provides-extra
-Dynamic: requires-dist
-Dynamic: requires-python
-Dynamic: summary
 
 # Neo4j Graph Data Science Client
 

graphdatascience-1.18/pyproject.toml

@@ -0,0 +1,175 @@
+[build-system]
+requires = ["setuptools>=42", "wheel"]
+build-backend = "setuptools.build_meta"
+
+[project]
+name = "graphdatascience"
+description = "A Python client for the Neo4j Graph Data Science (GDS) library"
+dynamic = ["version"]
+authors = [{ name = "Neo4j", email = "team-gds@neo4j.org" }]
+readme = "README.md"
+license = "Apache-2.0"
+classifiers = [
+    "Development Status :: 5 - Production/Stable",
+    "Intended Audience :: Developers",
+    "Intended Audience :: Science/Research",
+    "Operating System :: OS Independent",
+    "Programming Language :: Python :: 3",
+    "Programming Language :: Python :: 3 :: Only",
+    "Programming Language :: Python :: 3.10",
+    "Programming Language :: Python :: 3.11",
+    "Programming Language :: Python :: 3.12",
+    "Programming Language :: Python :: 3.13",
+    "Topic :: Database",
+    "Topic :: Scientific/Engineering",
+    "Topic :: Software Development",
+    "Typing :: Typed",
+]
+requires-python = ">=3.10"
+dependencies = [
+    "multimethod >= 1.0, < 3.0",
+    "neo4j >= 4.4.12, < 7.0",
+    "numpy < 2.4",
+    "pandas >= 1.0, < 3.0",
+    "pyarrow >= 17.0, < 23.0",
+    "textdistance >= 4.0, < 5.0",
+    "tqdm >= 4.0, < 5.0",
+    "typing-extensions >= 4.0, < 5.0",
+    "requests",
+    "tenacity >= 9.0",
+    "pydantic >= 2.11",
+]
+
+[project.urls]
+Homepage = "https://neo4j.com/product/graph-data-science/"
+Documentation = "https://neo4j.com/docs/graph-data-science-client/current/"
+Repository = "https://github.com/neo4j/graph-data-science-client"
+BugTracker = "https://github.com/neo4j/graph-data-science-client/issues"
+
+[project.optional-dependencies]
+networkx = ["networkx >= 2.0, < 4.0"]
+ogb = ["ogb >= 1.0, < 2.0"]
+rust-ext = ["neo4j-rust-ext >= 4.4.12, < 7.0"]
+
+
+[dependency-groups]
+# To install all development dependencies,
+# run `pip install --group dev -e .` inside repository root folder.#
+dev = [
+    "tox>4.24",
+    { include-group = "dev-base" },
+    { include-group = "test" },
+    { include-group = "docs-ci" },
+]
+
+
+dev-base = [
+    "ruff == 0.11.7",
+    "mypy == 1.13.0",
+    "nbconvert == 7.16.4",
+    "pandas-stubs == 2.2.3.241009",
+    "tox == 4.30.2",
+    "types-setuptools == 75.8.0.20250110",
+    "enum-tools[sphinx] == 0.12.0",
+    "autodoc_pydantic",
+    "types-requests",
+    "types-tqdm",
+    "types-python-dateutil",
+    "python-dotenv==1.1.0",
+    "types-colorama",
+]
+test = [
+    "pytest == 8.3.3",
+    "requests_mock == 1.11.0",
+    "pytest_mock == 3.15.1",
+    "testcontainers >= 4.0, < 4.13.0",
+    "python-dateutil >= 2.9",
+]
+docs-ci = [
+    "sphinx == 7.3.7",
+    "enum-tools[sphinx] == 0.12.0",
+    "autodoc_pydantic",
+]
+notebook-base = [
+    "nbconvert==7.16.4",
+    "nbformat==5.10.4",
+    "nbclient==0.10.2",
+    "ipykernel==6.29.5",
+]
+notebook-aura-ci = [{ include-group = "notebook-base" }, "python-dotenv==1.1.0"]
+notebook-ci = [
+    { include-group = "notebook-base" },
+    "scipy == 1.14.0",
+    "torch==2.3.0",
+    "torch-scatter==2.1.1",
+    "torch-sparse==0.6.17",
+    "torch-geometric>=2.5.0",
+]
+
+publish = ["build", "twine"]
+
+
+[tool.setuptools]
+include-package-data = true
+zip-safe = false
+
+[tool.setuptools.packages.find]
+where = ["src"]
+include = ["graphdatascience*"]
+
+[tool.setuptools.dynamic]
+version = { attr = "graphdatascience.version.__version__" }
+
+[tool.setuptools.package-data]
+graphdatascience = ["py.typed", "resources/**/*.gzip"]
+
+[tool.ruff]
+line-length = 120
+
+include = ["*.py", "*.ipynb"]
+
+exclude = [
+    # default
+
+    ".bzr",
+    ".direnv",
+    ".eggs",
+    ".git",
+    ".git-rewrite",
+    ".hg",
+    ".ipynb_checkpoints",
+    ".mypy_cache",
+    ".nox",
+    ".pants.d",
+    ".pyenv",
+    ".pytest_cache",
+    ".pytype",
+    ".ruff_cache",
+    ".svn",
+    ".tox",
+    ".venv",
+    ".vscode",
+    "__pypackages__",
+    "_build",
+    "buck-out",
+    "build",
+    "dist",
+    "node_modules",
+    "site-packages",
+    "venv",
+
+    # custom
+    "venv*",
+    ".venv*",
+    ".tox",
+    "build",
+]
+
+[tool.ruff.lint]
+select = [
+    "E4",
+    "E7",
+    "E9", # default pycodestyle rules
+    "F",  # flake8
+    "I",  # isort
+]

{graphdatascience-1.17 → graphdatascience-1.18/src}/graphdatascience/algo/algo_proc_runner.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 from abc import ABC
 from typing import Any
 
@@ -19,7 +21,7 @@ class AlgoProcRunner(IllegalAttrChecker, ABC):
         return self._query_runner.call_procedure(endpoint=self._namespace, params=params, logging=with_logging)
 
     @graph_type_check
-    def estimate(self, G: Graph, **config: Any) -> "Series[Any]":
+    def estimate(self, G: Graph, **config: Any) -> Series[Any]:
         self._namespace += "." + "estimate"
         return self._run_procedure(G, config, with_logging=False).squeeze()  # type: ignore
 
@@ -30,13 +32,13 @@ class StreamModeRunner(AlgoProcRunner):
 
 
 class StandardModeRunner(AlgoProcRunner):
-    def __call__(self, G: Graph, **config: Any) -> "Series[Any]":
+    def __call__(self, G: Graph, **config: Any) -> Series[Any]:
         return self._run_procedure(G, config).squeeze()  # type: ignore
 
 
 class GraphSageRunner(AlgoProcRunner):
     @graph_type_check
-    def __call__(self, G: Graph, **config: Any) -> tuple[GraphSageModel, "Series[Any]"]:
+    def __call__(self, G: Graph, **config: Any) -> tuple[GraphSageModel, Series[Any]]:
         result = self._run_procedure(G, config).squeeze()
         model_name = result["modelInfo"]["modelName"]
 

{graphdatascience-1.17 → graphdatascience-1.18/src}/graphdatascience/arrow_client/authenticated_flight_client.py

@@ -2,13 +2,17 @@ from __future__ import annotations
 
 import json
 import logging
+import platform
 from dataclasses import dataclass
-from typing import Any, Iterator, Optional, Tuple, Union
+from types import TracebackType
+from typing import Any, Iterator, Type
 
+import certifi
+from pyarrow import Schema, flight
 from pyarrow import __version__ as arrow_version
-from pyarrow import flight
 from pyarrow._flight import (
     Action,
+    ActionType,
     FlightInternalError,
     FlightStreamReader,
     FlightTimedOutError,
@@ -16,13 +20,11 @@ from pyarrow._flight import (
     Result,
     Ticket,
 )
-from tenacity import retry, retry_any, retry_if_exception_type, stop_after_attempt, stop_after_delay, wait_exponential
 
 from graphdatascience.arrow_client.arrow_authentication import ArrowAuthentication
 from graphdatascience.arrow_client.arrow_info import ArrowInfo
-from graphdatascience.retry_utils.retry_config import RetryConfig
+from graphdatascience.retry_utils.retry_config import ExponentialWaitConfig, RetryConfigV2, StopConfig
 
-from ..retry_utils.retry_utils import before_log
 from ..version import __version__
 from .middleware.auth_middleware import AuthFactory, AuthMiddleware
 from .middleware.user_agent_middleware import UserAgentFactory
@@ -32,12 +34,12 @@ class AuthenticatedArrowClient:
     @staticmethod
     def create(
         arrow_info: ArrowInfo,
-        auth: Optional[ArrowAuthentication] = None,
+        auth: ArrowAuthentication | None = None,
         encrypted: bool = False,
-        arrow_client_options: Optional[dict[str, Any]] = None,
-        connection_string_override: Optional[str] = None,
-        retry_config: Optional[RetryConfig] = None,
-        advertised_listen_address: Optional[Tuple[str, int]] = None,
+        arrow_client_options: dict[str, Any] | None = None,
+        connection_string_override: str | None = None,
+        retry_config: RetryConfigV2 | None = None,
+        advertised_listen_address: tuple[str, int] | None = None,
     ) -> AuthenticatedArrowClient:
         connection_string: str
         if connection_string_override is not None:
@@ -48,14 +50,14 @@ class AuthenticatedArrowClient:
         host, port = connection_string.split(":")
 
         if retry_config is None:
-            retry_config = RetryConfig(
-                retry=retry_any(
-                    retry_if_exception_type(FlightTimedOutError),
-                    retry_if_exception_type(FlightUnavailableError),
-                    retry_if_exception_type(FlightInternalError),
-                ),
-                stop=(stop_after_delay(10) | stop_after_attempt(5)),
-                wait=wait_exponential(multiplier=1, min=1, max=10),
+            retry_config = RetryConfigV2(
+                retryable_exceptions=[
+                    FlightTimedOutError,
+                    FlightUnavailableError,
+                    FlightInternalError,
+                ],
+                stop_config=StopConfig(after_delay=10, after_attempt=5),
+                wait_config=ExponentialWaitConfig(multiplier=1, min=1, max=10),
             )
 
         return AuthenticatedArrowClient(
@@ -71,15 +73,15 @@ class AuthenticatedArrowClient:
     def __init__(
         self,
         host: str,
-        retry_config: RetryConfig,
+        retry_config: RetryConfigV2,
         port: int = 8491,
-        auth: Optional[ArrowAuthentication] = None,
+        auth: ArrowAuthentication | None = None,
         encrypted: bool = False,
-        arrow_client_options: Optional[dict[str, Any]] = None,
-        user_agent: Optional[str] = None,
-        advertised_listen_address: Optional[Tuple[str, int]] = None,
+        arrow_client_options: dict[str, Any] | None = None,
+        user_agent: str | None = None,
+        advertised_listen_address: tuple[str, int] | None = None,
     ):
-        """Creates a new GdsArrowClient instance.
+        """Creates a new AuthenticatedArrowClient instance.
 
         Parameters
         ----------
@@ -87,17 +89,17 @@ class AuthenticatedArrowClient:
             The host address of the GDS Arrow server
         port: int
             The host port of the GDS Arrow server (default is 8491)
-        auth: Optional[ArrowAuthentication]
+        auth: ArrowAuthentication | None
             An implementation of ArrowAuthentication providing a pair to be used for basic authentication
         encrypted: bool
             A flag that indicates whether the connection should be encrypted (default is False)
-        arrow_client_options: Optional[dict[str, Any]]
+        arrow_client_options: dict[str, Any] | None
             Additional options to be passed to the Arrow Flight client.
-        user_agent: Optional[str]
+        user_agent: str | None
             The user agent string to use for the connection. (default is `neo4j-graphdatascience-v[VERSION] pyarrow-v[PYARROW_VERSION])
-        retry_config: Optional[RetryConfig]
+        retry_config: RetryConfig | None
             The retry configuration to use for the Arrow requests send by the client.
-        advertised_listen_address: Optional[Tuple[str, int]]
+        advertised_listen_address: tuple[str, int] | None
             The advertised listen address of the GDS Arrow server. This will be used by remote projection and writeback operations.
         """
         self._host = host
@@ -106,7 +108,6 @@ class AuthenticatedArrowClient:
         self._encrypted = encrypted
         self._arrow_client_options = arrow_client_options
         self._user_agent = user_agent
-        self._retry_config = retry_config
         self._logger = logging.getLogger("gds_arrow_client")
         self._retry_config = retry_config
         if auth:
@@ -114,7 +115,7 @@ class AuthenticatedArrowClient:
             self._auth_middleware = AuthMiddleware(auth)
         self.advertised_listen_address = advertised_listen_address
 
-        self._flight_client = self._instantiate_flight_client()
+        self._flight_client: flight.FlightClient = self._instantiate_flight_client()
 
     def connection_info(self) -> ConnectionInfo:
         """
@@ -122,7 +123,7 @@ class AuthenticatedArrowClient:
 
         Returns
         -------
-        tuple[str, int]
+        ConnectionInfo
             the host and port of the GDS Arrow server
         """
         return ConnectionInfo(self._host, self._port, self._encrypted)
@@ -142,23 +143,17 @@ class AuthenticatedArrowClient:
         h, p = self.advertised_listen_address
         return ConnectionInfo(h, p, self._encrypted)
 
-    def request_token(self) -> Optional[str]:
+    def request_token(self) -> str | None:
         """
         Requests a token from the server and returns it.
 
         Returns
         -------
-        Optional[str]
+        str | None
             a token from the server and returns it.
         """
 
-        @retry(
-            reraise=True,
-            before=before_log("Request token", self._logger, logging.DEBUG),
-            retry=self._retry_config.retry,
-            stop=self._retry_config.stop,
-            wait=self._retry_config.wait,
-        )
+        @self._retry_config.decorator(operation_name="Request token", logger=self._logger)
         def auth_with_retry() -> None:
             client = self._flight_client
             if self._auth:
@@ -174,24 +169,47 @@ class AuthenticatedArrowClient:
     def get_stream(self, ticket: Ticket) -> FlightStreamReader:
         return self._flight_client.do_get(ticket)
 
-    def do_action(self, endpoint: str, payload: Union[bytes, dict[str, Any]]) -> Iterator[Result]:
+    def do_action(self, endpoint: str, payload: bytes | dict[str, Any]) -> Iterator[Result]:
         payload_bytes = payload if isinstance(payload, bytes) else json.dumps(payload).encode("utf-8")
 
         return self._flight_client.do_action(Action(endpoint, payload_bytes))  # type: ignore
 
-    def do_action_with_retry(self, endpoint: str, payload: Union[bytes, dict[str, Any]]) -> Iterator[Result]:
-        @retry(
-            reraise=True,
-            before=before_log("Send action", self._logger, logging.DEBUG),
-            retry=self._retry_config.retry,
-            stop=self._retry_config.stop,
-            wait=self._retry_config.wait,
-        )
-        def run_with_retry() -> Iterator[Result]:
-            return self.do_action(endpoint, payload)
+    def do_action_with_retry(self, endpoint: str, payload: bytes | dict[str, Any]) -> list[Result]:
+        @self._retry_config.decorator(operation_name="Send action", logger=self._logger)
+        def run_with_retry() -> list[Result]:
+            # the Flight response error code is only checked on iterator consumption
+            # we eagerly collect iterator here to trigger retry in case of an error
+            return list(self.do_action(endpoint, payload))
+
+        return run_with_retry()
+
+    def list_actions(self) -> set[ActionType]:
+        return self._flight_client.list_actions()  # type: ignore
+
+    def do_put_with_retry(
+        self, descriptor: flight.FlightDescriptor, schema: Schema
+    ) -> tuple[flight.FlightStreamWriter, flight.FlightMetadataReader]:
+        @self._retry_config.decorator(operation_name="Do put", logger=self._logger)
+        def run_with_retry() -> tuple[flight.FlightStreamWriter, flight.FlightMetadataReader]:
+            return self._flight_client.do_put(descriptor, schema)  # type: ignore
 
         return run_with_retry()
 
+    def __enter__(self) -> AuthenticatedArrowClient:
+        return self
+
+    def __exit__(
+        self,
+        exception_type: Type[BaseException] | None,
+        exception_value: BaseException | None,
+        traceback: TracebackType | None,
+    ) -> None:
+        self.close()
+
+    def close(self) -> None:
+        if self._flight_client:
+            self._flight_client.close()
+
     def _instantiate_flight_client(self) -> flight.FlightClient:
         location = (
             flight.Location.for_grpc_tls(self._host, self._port)
@@ -199,6 +217,12 @@ class AuthenticatedArrowClient:
             else flight.Location.for_grpc_tcp(self._host, self._port)
         )
         client_options: dict[str, Any] = (self._arrow_client_options or {}).copy()
+
+        # We need to specify the system root certificates on Windows
+        if platform.system() == "Windows":
+            if not client_options["tls_root_certs"]:
+                client_options["tls_root_certs"] = certifi.contents()
+
         if self._auth:
             user_agent = f"neo4j-graphdatascience-v{__version__} pyarrow-v{arrow_version}"
             if self._user_agent:
@@ -211,6 +235,17 @@ class AuthenticatedArrowClient:
 
         return flight.FlightClient(location, **client_options)
 
+    def __getstate__(self) -> dict[str, Any]:
+        state = self.__dict__.copy()
+        # Remove the FlightClient as it isn't serializable
+        if "_flight_client" in state:
+            del state["_flight_client"]
+        return state
+
+    def __setstate__(self, state: dict[str, Any]) -> None:
+        self.__dict__.update(state)
+        self._flight_client = self._instantiate_flight_client()
+
 
 @dataclass
 class ConnectionInfo:

{graphdatascience-1.17 → graphdatascience-1.18/src}/graphdatascience/arrow_client/middleware/auth_middleware.py

@@ -2,7 +2,7 @@ from __future__ import annotations
 
 import base64
 import time
-from typing import Any, Optional
+from typing import Any
 
 from pyarrow._flight import ClientMiddleware, ClientMiddlewareFactory
 
@@ -22,10 +22,10 @@ class AuthMiddleware(ClientMiddleware):  # type: ignore
     def __init__(self, auth: ArrowAuthentication, *args: Any, **kwargs: Any) -> None:
         super().__init__(*args, **kwargs)
         self._auth = auth
-        self._token: Optional[str] = None
+        self._token: str | None = None
         self._token_timestamp = 0
 
-    def token(self) -> Optional[str]:
+    def token(self) -> str | None:
         # check whether the token is older than 10 minutes. If so, reset it.
         if self._token and int(time.time()) - self._token_timestamp > 600:
             self._token = None

{graphdatascience-1.17/graphdatascience/arrow_client/v2 → graphdatascience-1.18/src/graphdatascience/arrow_client/v1}/data_mapper_utils.py

@@ -1,10 +1,10 @@
 import json
-from typing import Any, Iterator
+from typing import Any
 
 from pyarrow._flight import Result
 
 
-def deserialize_single(input_stream: Iterator[Result]) -> dict[str, Any]:
+def deserialize_single(input_stream: list[Result]) -> dict[str, Any]:
     rows = deserialize(input_stream)
     if len(rows) != 1:
         raise ValueError(f"Expected exactly one result, got {len(rows)}")
@@ -12,8 +12,8 @@ def deserialize_single(input_stream: Iterator[Result]) -> dict[str, Any]:
     return rows[0]
 
 
-def deserialize(input_stream: Iterator[Result]) -> list[dict[str, Any]]:
+def deserialize(input_stream: list[Result]) -> list[dict[str, Any]]:
     def deserialize_row(row: Result):  # type:ignore
         return json.loads(row.body.to_pybytes().decode())
 
-    return [deserialize_row(row) for row in list(input_stream)]
+    return [deserialize_row(row) for row in input_stream]