graphdatascience 1.18a1__tar.gz → 1.19a1__tar.gz

graphdatascience-1.19a1/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.18a1/graphdatascience.egg-info → graphdatascience-1.19a1}/PKG-INFO

@@ -1,18 +1,16 @@
 Metadata-Version: 2.4
 Name: graphdatascience
-Version: 1.18a1
+Version: 1.19a1
 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
@@ -31,32 +29,20 @@ Requires-Dist: multimethod<3.0,>=1.0
 Requires-Dist: neo4j<7.0,>=4.4.12
 Requires-Dist: numpy<2.4
 Requires-Dist: pandas<3.0,>=1.0
-Requires-Dist: pyarrow<22.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.19a1/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.19a1/src/graphdatascience/arrow_client/arrow_client_options_util.py

@@ -0,0 +1,13 @@
+from typing import Any
+
+TLS_ROOT_CERTS_OPTION = "tls_root_certs"
+
+
+def disable_server_verification(arrow_client_options: dict[str, Any]) -> dict[str, Any]:
+    arrow_client_options["disable_server_verification"] = True
+    return arrow_client_options
+
+
+def set_tls_root_certs(arrow_client_options: dict[str, Any], tls_root_certs: bytes | str) -> dict[str, Any]:
+    arrow_client_options[TLS_ROOT_CERTS_OPTION] = tls_root_certs
+    return arrow_client_options

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

@@ -2,11 +2,14 @@ from __future__ import annotations
 
 import json
 import logging
+import platform
 from dataclasses import dataclass
-from typing import Any, Iterator
+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,
@@ -17,14 +20,13 @@ 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 .arrow_client_options_util import TLS_ROOT_CERTS_OPTION, set_tls_root_certs
 from .middleware.auth_middleware import AuthFactory, AuthMiddleware
 from .middleware.user_agent_middleware import UserAgentFactory
 
@@ -37,7 +39,7 @@ class AuthenticatedArrowClient:
         encrypted: bool = False,
         arrow_client_options: dict[str, Any] | None = None,
         connection_string_override: str | None = None,
-        retry_config: RetryConfig | None = None,
+        retry_config: RetryConfigV2 | None = None,
         advertised_listen_address: tuple[str, int] | None = None,
     ) -> AuthenticatedArrowClient:
         connection_string: str
@@ -49,14 +51,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(
@@ -72,7 +74,7 @@ class AuthenticatedArrowClient:
     def __init__(
         self,
         host: str,
-        retry_config: RetryConfig,
+        retry_config: RetryConfigV2,
         port: int = 8491,
         auth: ArrowAuthentication | None = None,
         encrypted: bool = False,
@@ -80,7 +82,7 @@ class AuthenticatedArrowClient:
         user_agent: str | None = None,
         advertised_listen_address: tuple[str, int] | None = None,
     ):
-        """Creates a new GdsArrowClient instance.
+        """Creates a new AuthenticatedArrowClient instance.
 
         Parameters
         ----------
@@ -107,7 +109,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:
@@ -115,7 +116,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:
         """
@@ -123,7 +124,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)
@@ -153,13 +154,7 @@ class AuthenticatedArrowClient:
             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:
@@ -181,13 +176,7 @@ class AuthenticatedArrowClient:
         return self._flight_client.do_action(Action(endpoint, payload_bytes))  # type: ignore
 
     def do_action_with_retry(self, endpoint: str, payload: bytes | dict[str, Any]) -> list[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,
-        )
+        @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
@@ -198,6 +187,30 @@ class AuthenticatedArrowClient:
     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)
@@ -205,6 +218,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 TLS_ROOT_CERTS_OPTION not in client_options:
+                set_tls_root_certs(client_options, certifi.contents())
+
         if self._auth:
             user_agent = f"neo4j-graphdatascience-v{__version__} pyarrow-v{arrow_version}"
             if self._user_agent:
@@ -217,6 +236,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: