diff --git a/pyproject.toml b/pyproject.toml index 646c7aab..0e649fa7 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -117,5 +117,9 @@ asyncio_mode = "auto" [tool.mypy] [[tool.mypy.overrides]] -module = "pandas.*" +module = [ + "pandas.*", + "neo4j._codec.packstream._rust", + "neo4j._codec.packstream._rust.*", +] ignore_missing_imports = true diff --git a/src/neo4j/_codec/packstream/_common.py b/src/neo4j/_codec/packstream/_common.py index 84403de7..30b97a0d 100644 --- a/src/neo4j/_codec/packstream/_common.py +++ b/src/neo4j/_codec/packstream/_common.py @@ -16,29 +16,7 @@ # limitations under the License. -class Structure: - - def __init__(self, tag, *fields): - self.tag = tag - self.fields = list(fields) - - def __repr__(self): - return "Structure[0x%02X](%s)" % (ord(self.tag), ", ".join(map(repr, self.fields))) - - def __eq__(self, other): - try: - return self.tag == other.tag and self.fields == other.fields - except AttributeError: - return False - - def __ne__(self, other): - return not self.__eq__(other) - - def __len__(self): - return len(self.fields) - - def __getitem__(self, key): - return self.fields[key] - - def __setitem__(self, key, value): - self.fields[key] = value +try: + from ._rust import Structure +except ImportError: + from ._python import Structure diff --git a/src/neo4j/_codec/packstream/_python/__init__.py b/src/neo4j/_codec/packstream/_python/__init__.py new file mode 100644 index 00000000..ba0188b0 --- /dev/null +++ b/src/neo4j/_codec/packstream/_python/__init__.py @@ -0,0 +1,24 @@ +# Copyright (c) "Neo4j" +# Neo4j Sweden AB [https://neo4j.com] +# +# This file is part of Neo4j. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + + +from ._common import Structure + + +__all__ = [ + "Structure", +] diff --git a/src/neo4j/_codec/packstream/_python/_common.py b/src/neo4j/_codec/packstream/_python/_common.py new file mode 100644 index 00000000..0bca00cd --- /dev/null +++ b/src/neo4j/_codec/packstream/_python/_common.py @@ -0,0 +1,46 @@ +# Copyright (c) "Neo4j" +# Neo4j Sweden AB [https://neo4j.com] +# +# This file is part of Neo4j. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + + +class Structure: + + def __init__(self, tag, *fields): + self.tag = tag + self.fields = list(fields) + + def __repr__(self): + return "Structure[0x%02X](%s)" % ( + ord(self.tag), ", ".join(map(repr, self.fields)) + ) + + def __eq__(self, other): + try: + return self.tag == other.tag and self.fields == other.fields + except AttributeError: + return False + + def __ne__(self, other): + return not self.__eq__(other) + + def __len__(self): + return len(self.fields) + + def __getitem__(self, key): + return self.fields[key] + + def __setitem__(self, key, value): + self.fields[key] = value diff --git a/src/neo4j/_codec/packstream/v1/__init__.py b/src/neo4j/_codec/packstream/v1/__init__.py index 360cd25b..c28f069b 100644 --- a/src/neo4j/_codec/packstream/v1/__init__.py +++ b/src/neo4j/_codec/packstream/v1/__init__.py @@ -16,7 +16,6 @@ # limitations under the License. -import typing as t from codecs import decode from contextlib import contextmanager from struct import ( @@ -24,38 +23,19 @@ unpack as struct_unpack, ) -from ...._optional_deps import ( - np, - pd, -) from ...hydration import DehydrationHooks from .._common import Structure +from .types import * -NONE_VALUES: t.Tuple = (None,) -TRUE_VALUES: t.Tuple = (True,) -FALSE_VALUES: t.Tuple = (False,) -INT_TYPES: t.Tuple[t.Type, ...] = (int,) -FLOAT_TYPES: t.Tuple[t.Type, ...] = (float,) -# we can't put tuple here because spatial types subclass tuple, -# and we don't want to treat them as sequences -SEQUENCE_TYPES: t.Tuple[t.Type, ...] = (list,) -MAPPING_TYPES: t.Tuple[t.Type, ...] = (dict,) -BYTES_TYPES: t.Tuple[t.Type, ...] = (bytes, bytearray) - - -if np is not None: - TRUE_VALUES = (*TRUE_VALUES, np.bool_(True)) - FALSE_VALUES = (*FALSE_VALUES, np.bool_(False)) - INT_TYPES = (*INT_TYPES, np.integer) - FLOAT_TYPES = (*FLOAT_TYPES, np.floating) - SEQUENCE_TYPES = (*SEQUENCE_TYPES, np.ndarray) - -if pd is not None: - NONE_VALUES = (*NONE_VALUES, pd.NA) - SEQUENCE_TYPES = (*SEQUENCE_TYPES, pd.Series, pd.Categorical, - pd.core.arrays.ExtensionArray) - MAPPING_TYPES = (*MAPPING_TYPES, pd.DataFrame) +try: + from .._rust.v1 import ( + pack as _rust_pack, + unpack as _rust_unpack, + ) +except ImportError: + _rust_pack = None + _rust_unpack = None PACKED_UINT_8 = [struct_pack(">B", value) for value in range(0x100)] @@ -74,12 +54,17 @@ def __init__(self, stream): self.stream = stream self._write = self.stream.write - def _pack_raw(self, data): - self._write(data) - def pack(self, data, dehydration_hooks=None): - self._pack(data, - dehydration_hooks=self._inject_hooks(dehydration_hooks)) + dehydration_hooks = self._inject_hooks(dehydration_hooks) + self._pack(data, dehydration_hooks=dehydration_hooks) + + if _rust_pack: + def _pack(self, data, dehydration_hooks=None): + data = _rust_pack(data, dehydration_hooks) + self._write(data) + else: + def _pack(self, data, dehydration_hooks=None): + self._py_pack(data, dehydration_hooks) @classmethod def _inject_hooks(cls, dehydration_hooks=None): @@ -93,8 +78,7 @@ def _inject_hooks(cls, dehydration_hooks=None): subtypes={} ) - - def _pack(self, value, dehydration_hooks=None): + def _py_pack(self, value, dehydration_hooks=None): write = self._write # None @@ -136,18 +120,18 @@ def _pack(self, value, dehydration_hooks=None): elif isinstance(value, str): encoded = value.encode("utf-8") self._pack_string_header(len(encoded)) - self._pack_raw(encoded) + self._write(encoded) # Bytes elif isinstance(value, BYTES_TYPES): self._pack_bytes_header(len(value)) - self._pack_raw(value) + self._write(value) # List elif isinstance(value, SEQUENCE_TYPES): self._pack_list_header(len(value)) for item in value: - self._pack(item, dehydration_hooks) + self._py_pack(item, dehydration_hooks) # Map elif isinstance(value, MAPPING_TYPES): @@ -157,8 +141,8 @@ def _pack(self, value, dehydration_hooks=None): raise TypeError( "Map keys must be strings, not {}".format(type(key)) ) - self._pack(key, dehydration_hooks) - self._pack(item, dehydration_hooks) + self._py_pack(key, dehydration_hooks) + self._py_pack(item, dehydration_hooks) # Structure elif isinstance(value, Structure): @@ -169,7 +153,7 @@ def _pack(self, value, dehydration_hooks=None): if dehydration_hooks: transformer = dehydration_hooks.get_transformer(value) if transformer is not None: - self._pack(transformer(value), dehydration_hooks) + self._py_pack(transformer(value), dehydration_hooks) return raise ValueError("Values of type %s are not supported" % type(value)) @@ -298,11 +282,16 @@ def read(self, n=1): def read_u8(self): return self.unpackable.read_u8() - def unpack(self, hydration_hooks=None): - value = self._unpack(hydration_hooks=hydration_hooks) - if hydration_hooks and type(value) in hydration_hooks: - return hydration_hooks[type(value)](value) - return value + if _rust_unpack: + def unpack(self, hydration_hooks=None): + value, i = _rust_unpack( + self.unpackable.data, self.unpackable.p, hydration_hooks + ) + self.unpackable.p = i + return value + else: + def unpack(self, hydration_hooks=None): + return self._unpack(hydration_hooks=hydration_hooks) def _unpack(self, hydration_hooks=None): marker = self.read_u8() @@ -384,8 +373,13 @@ def _unpack(self, hydration_hooks=None): size, tag = self._unpack_structure_header(marker) value = Structure(tag, *([None] * size)) for i in range(len(value)): - value[i] = self.unpack(hydration_hooks=hydration_hooks) - return value + value[i] = self._unpack(hydration_hooks=hydration_hooks) + if not hydration_hooks: + return value + hydration_hook = hydration_hooks.get(type(value)) + if not hydration_hook: + return value + return hydration_hook(value) else: raise ValueError("Unknown PackStream marker %02X" % marker) @@ -397,22 +391,22 @@ def _unpack_list_items(self, marker, hydration_hooks=None): if size == 0: return elif size == 1: - yield self.unpack(hydration_hooks=hydration_hooks) + yield self._unpack(hydration_hooks=hydration_hooks) else: for _ in range(size): - yield self.unpack(hydration_hooks=hydration_hooks) + yield self._unpack(hydration_hooks=hydration_hooks) elif marker == 0xD4: # LIST_8: size, = struct_unpack(">B", self.read(1)) for _ in range(size): - yield self.unpack(hydration_hooks=hydration_hooks) + yield self._unpack(hydration_hooks=hydration_hooks) elif marker == 0xD5: # LIST_16: size, = struct_unpack(">H", self.read(2)) for _ in range(size): - yield self.unpack(hydration_hooks=hydration_hooks) + yield self._unpack(hydration_hooks=hydration_hooks) elif marker == 0xD6: # LIST_32: size, = struct_unpack(">I", self.read(4)) for _ in range(size): - yield self.unpack(hydration_hooks=hydration_hooks) + yield self._unpack(hydration_hooks=hydration_hooks) else: return @@ -426,29 +420,29 @@ def _unpack_map(self, marker, hydration_hooks=None): size = marker & 0x0F value = {} for _ in range(size): - key = self.unpack(hydration_hooks=hydration_hooks) - value[key] = self.unpack(hydration_hooks=hydration_hooks) + key = self._unpack(hydration_hooks=hydration_hooks) + value[key] = self._unpack(hydration_hooks=hydration_hooks) return value elif marker == 0xD8: # MAP_8: size, = struct_unpack(">B", self.read(1)) value = {} for _ in range(size): - key = self.unpack(hydration_hooks=hydration_hooks) - value[key] = self.unpack(hydration_hooks=hydration_hooks) + key = self._unpack(hydration_hooks=hydration_hooks) + value[key] = self._unpack(hydration_hooks=hydration_hooks) return value elif marker == 0xD9: # MAP_16: size, = struct_unpack(">H", self.read(2)) value = {} for _ in range(size): - key = self.unpack(hydration_hooks=hydration_hooks) - value[key] = self.unpack(hydration_hooks=hydration_hooks) + key = self._unpack(hydration_hooks=hydration_hooks) + value[key] = self._unpack(hydration_hooks=hydration_hooks) return value elif marker == 0xDA: # MAP_32: size, = struct_unpack(">I", self.read(4)) value = {} for _ in range(size): - key = self.unpack(hydration_hooks=hydration_hooks) - value[key] = self.unpack(hydration_hooks=hydration_hooks) + key = self._unpack(hydration_hooks=hydration_hooks) + value[key] = self._unpack(hydration_hooks=hydration_hooks) return value else: return None diff --git a/src/neo4j/_codec/packstream/v1/types.py b/src/neo4j/_codec/packstream/v1/types.py new file mode 100644 index 00000000..591b8eba --- /dev/null +++ b/src/neo4j/_codec/packstream/v1/types.py @@ -0,0 +1,62 @@ +# Copyright (c) "Neo4j" +# Neo4j Sweden AB [https://neo4j.com] +# +# This file is part of Neo4j. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + + +import typing as t + +from ...._optional_deps import ( + np, + pd, +) + + +NONE_VALUES: t.Tuple = (None,) +TRUE_VALUES: t.Tuple = (True,) +FALSE_VALUES: t.Tuple = (False,) +INT_TYPES: t.Tuple[t.Type, ...] = (int,) +FLOAT_TYPES: t.Tuple[t.Type, ...] = (float,) +# we can't put tuple here because spatial types subclass tuple, +# and we don't want to treat them as sequences +SEQUENCE_TYPES: t.Tuple[t.Type, ...] = (list,) +MAPPING_TYPES: t.Tuple[t.Type, ...] = (dict,) +BYTES_TYPES: t.Tuple[t.Type, ...] = (bytes, bytearray) + + +if np is not None: + TRUE_VALUES = (*TRUE_VALUES, np.bool_(True)) + FALSE_VALUES = (*FALSE_VALUES, np.bool_(False)) + INT_TYPES = (*INT_TYPES, np.integer) + FLOAT_TYPES = (*FLOAT_TYPES, np.floating) + SEQUENCE_TYPES = (*SEQUENCE_TYPES, np.ndarray) + +if pd is not None: + NONE_VALUES = (*NONE_VALUES, pd.NA) + SEQUENCE_TYPES = (*SEQUENCE_TYPES, pd.Series, pd.Categorical, + pd.core.arrays.ExtensionArray) + MAPPING_TYPES = (*MAPPING_TYPES, pd.DataFrame) + + +__all__ = [ + "NONE_VALUES", + "TRUE_VALUES", + "FALSE_VALUES", + "INT_TYPES", + "FLOAT_TYPES", + "SEQUENCE_TYPES", + "MAPPING_TYPES", + "BYTES_TYPES", +] diff --git a/testkit/Dockerfile b/testkit/Dockerfile index e407b453..a5c17e94 100644 --- a/testkit/Dockerfile +++ b/testkit/Dockerfile @@ -67,4 +67,4 @@ RUN apt update && \ wget https://apache.jfrog.io/artifactory/arrow/${distro_name}/apache-arrow-apt-source-latest-${code_name}.deb && \ apt install -y -V ./apache-arrow-apt-source-latest-${code_name}.deb && \ apt update && \ - apt install -y -V libarrow-dev # For C++ \ + apt install -y -V libarrow-dev # For C++ diff --git a/testkit/_common.py b/testkit/_common.py index 76ad479f..2d93e635 100644 --- a/testkit/_common.py +++ b/testkit/_common.py @@ -1,3 +1,21 @@ +# Copyright (c) "Neo4j" +# Neo4j Sweden AB [https://neo4j.com] +# +# This file is part of Neo4j. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + + import os import re import subprocess diff --git a/testkit/backend.py b/testkit/backend.py index 702ea8a5..59b4690e 100644 --- a/testkit/backend.py +++ b/testkit/backend.py @@ -1,5 +1,3 @@ -#!/usr/bin/env python - # Copyright (c) "Neo4j" # Neo4j Sweden AB [https://neo4j.com] # diff --git a/testkit/build.py b/testkit/build.py index 3ac90a2f..b13ee0d6 100644 --- a/testkit/build.py +++ b/testkit/build.py @@ -1,5 +1,3 @@ -#!/usr/bin/env python - # Copyright (c) "Neo4j" # Neo4j Sweden AB [https://neo4j.com] # diff --git a/testkit/integration.py b/testkit/integration.py index edb08bd7..67dd596d 100644 --- a/testkit/integration.py +++ b/testkit/integration.py @@ -1,5 +1,3 @@ -#!/usr/bin/env python - # Copyright (c) "Neo4j" # Neo4j Sweden AB [https://neo4j.com] # diff --git a/testkit/stress.py b/testkit/stress.py index 2b4fdc57..3828fe34 100644 --- a/testkit/stress.py +++ b/testkit/stress.py @@ -1,5 +1,3 @@ -#!/usr/bin/env python - # Copyright (c) "Neo4j" # Neo4j Sweden AB [https://neo4j.com] # diff --git a/testkit/unittests.py b/testkit/unittests.py index 0b714d3d..4cd014b7 100644 --- a/testkit/unittests.py +++ b/testkit/unittests.py @@ -1,5 +1,3 @@ -#!/usr/bin/env python - # Copyright (c) "Neo4j" # Neo4j Sweden AB [https://neo4j.com] # diff --git a/testkitbackend/_async/backend.py b/testkitbackend/_async/backend.py index 89d35578..2d4a7cb7 100644 --- a/testkitbackend/_async/backend.py +++ b/testkitbackend/_async/backend.py @@ -28,6 +28,7 @@ ) from pathlib import Path +import neo4j from neo4j._exceptions import BoltError from neo4j.exceptions import ( DriverError, @@ -45,7 +46,7 @@ TESTKIT_BACKEND_PATH = Path(__file__).absolute().resolve().parents[1] -DRIVER_PATH = TESTKIT_BACKEND_PATH.parent / "src" / "neo4j" +DRIVER_PATH = Path(neo4j.__path__[0]).absolute().resolve() class AsyncBackend: diff --git a/testkitbackend/_sync/backend.py b/testkitbackend/_sync/backend.py index a38aa1e0..0e88a0fa 100644 --- a/testkitbackend/_sync/backend.py +++ b/testkitbackend/_sync/backend.py @@ -28,6 +28,7 @@ ) from pathlib import Path +import neo4j from neo4j._exceptions import BoltError from neo4j.exceptions import ( DriverError, @@ -45,7 +46,7 @@ TESTKIT_BACKEND_PATH = Path(__file__).absolute().resolve().parents[1] -DRIVER_PATH = TESTKIT_BACKEND_PATH.parent / "src" / "neo4j" +DRIVER_PATH = Path(neo4j.__path__[0]).absolute().resolve() class Backend: diff --git a/tests/unit/common/codec/packstream/v1/test_packstream.py b/tests/unit/common/codec/packstream/v1/test_packstream.py index 7e5bd493..a649755a 100644 --- a/tests/unit/common/codec/packstream/v1/test_packstream.py +++ b/tests/unit/common/codec/packstream/v1/test_packstream.py @@ -52,11 +52,13 @@ def unpacker_with_buffer(): unpackable_buffer = Unpacker.new_unpackable_buffer() return Unpacker(unpackable_buffer), unpackable_buffer + def test_packable_buffer(packer_with_buffer): packer, packable_buffer = packer_with_buffer assert isinstance(packable_buffer, PackableBuffer) assert packable_buffer is packer.stream + def test_unpackable_buffer(unpacker_with_buffer): unpacker, unpackable_buffer = unpacker_with_buffer assert isinstance(unpackable_buffer, UnpackableBuffer) @@ -122,7 +124,6 @@ def np_float_overflow_as_error(request): np.seterr(**old_err) - @pytest.fixture(params=( int, np.int8, np.int16, np.int32, np.int64, np.longlong, @@ -547,7 +548,7 @@ def test_map_key_type(self, packer_with_buffer, map_, exc_type): # maps must have string keys packer, packable_buffer = packer_with_buffer with pytest.raises(exc_type, match="strings"): - packer._pack(map_) + packer._py_pack(map_) def test_illegal_signature(self, assert_packable): with pytest.raises(ValueError):