pyetp 0.0.43__py3-none-any.whl → 0.0.44__py3-none-any.whl

pyetp/__init__.py

@@ -1,5 +1,5 @@
 from ._version import __version__
-from .client import ETPClient, ETPError, connect
+from .client import ETPClient, ETPError, connect, etp_connect, etp_persistent_connect
 from .uri import DataObjectURI, DataspaceURI
 
 __all__ = [
@@ -7,6 +7,8 @@ __all__ = [
     "ETPClient",
     "ETPError",
     "connect",
+    "etp_connect",
+    "etp_persistent_connect",
     "DataObjectURI",
     "DataspaceURI",
 ]

pyetp/_version.py

@@ -28,7 +28,7 @@ version_tuple: VERSION_TUPLE
 commit_id: COMMIT_ID
 __commit_id__: COMMIT_ID
 
-__version__ = version = '0.0.43'
-__version_tuple__ = version_tuple = (0, 0, 43)
+__version__ = version = '0.0.44'
+__version_tuple__ = version_tuple = (0, 0, 44)
 
 __commit_id__ = commit_id = None

pyetp/client.py

@@ -1,4 +1,5 @@
 import asyncio
+import contextlib
 import datetime
 import logging
 import sys
@@ -6,10 +7,13 @@ import typing as T
 import uuid
 import warnings
 from collections import defaultdict
+from collections.abc import AsyncGenerator
 from types import TracebackType
 
 import numpy as np
+import numpy.typing as npt
 import websockets
+import websockets.client
 from etpproto.connection import CommunicationProtocol, ConnectionType, ETPConnection
 from etpproto.messages import Message, MessageFlags
 from etptypes import ETPModel
@@ -143,6 +147,7 @@ from xtgeo import RegularSurface
 
 import resqml_objects.v201 as ro
 from pyetp import utils_arrays, utils_xml
+from pyetp._version import version
 from pyetp.config import SETTINGS
 from pyetp.uri import DataObjectURI, DataspaceURI
 from resqml_objects import parse_resqml_v201_object, serialize_resqml_v201_object
@@ -175,6 +180,14 @@ except ImportError:
         except asyncio.CancelledError as e:
             raise asyncio.TimeoutError(f"Timeout ({delay}s)") from e
 
+    TimeoutError = asyncio.TimeoutError
+
+try:
+    # Python >= 3.11
+    from typing import Self
+except ImportError:
+    Self = "ETPClient"
+
 
 class ETPError(Exception):
     def __init__(self, message: str, code: int):
@@ -192,6 +205,10 @@ class ETPError(Exception):
         return list(map(cls.from_proto, errors))
 
 
+class ReceiveWorkerExited(Exception):
+    pass
+
+
 def get_all_etp_protocol_classes():
     """Update protocol - all exception protocols are now per message"""
 
@@ -211,17 +228,31 @@ class ETPClient(ETPConnection):
     _recv_events: T.Dict[int, asyncio.Event]
     _recv_buffer: T.Dict[int, T.List[ETPModel]]
 
-    def __init__(self, ws: websockets.ClientConnection, timeout=10.0):
+    def __init__(
+        self,
+        ws: websockets.ClientConnection,
+        etp_timeout: float | None = 10.0,
+        max_message_size: float = 2**20,
+        application_name: str = "pyetp",
+        application_version: str = version,
+    ) -> None:
         super().__init__(connection_type=ConnectionType.CLIENT)
+
+        self.application_name = application_name
+        self.application_version = application_version
+
         self._recv_events = {}
         self._recv_buffer = defaultdict(lambda: list())  # type: ignore
         self.ws = ws
 
-        self.timeout = timeout
+        # Ensure a minimum timeout of 10 seconds.
+        self.etp_timeout = (
+            etp_timeout if etp_timeout is None or etp_timeout > 10.0 else 10.0
+        )
         self.client_info.endpoint_capabilities["MaxWebSocketMessagePayloadSize"] = (
-            SETTINGS.MaxWebSocketMessagePayloadSize
+            max_message_size
         )
-        self.__recvtask = asyncio.create_task(self.__recv__())
+        self.__recvtask = asyncio.create_task(self.__recv())
 
     #
     # client
@@ -242,25 +273,81 @@ class ETPClient(ETPConnection):
         # create future recv event
         self._recv_events[msg.header.message_id] = asyncio.Event()
 
+        tasks = []
         for msg_part in msg.encode_message_generator(self.max_size, self):
-            await self.ws.send(msg_part)
+            tasks.append(self.ws.send(msg_part))
+
+        await asyncio.gather(*tasks)
 
         return msg.header.message_id
 
     async def _recv(self, correlation_id: int) -> ETPModel:
         assert correlation_id in self._recv_events, (
-            "trying to recv response on non-existing message"
+            "Trying to receive a response on non-existing message"
         )
 
-        async with timeout(self.timeout):
-            await self._recv_events[correlation_id].wait()
+        def timeout_intervals(etp_timeout):
+            # Local function generating progressively longer timeout intervals.
+
+            # Use the timeout-interval generator from the Python websockets
+            # library.
+            backoff_generator = websockets.client.backoff(
+                initial_delay=5.0, min_delay=5.0, max_delay=20.0
+            )
+
+            # Check if we should never time out.
+            if etp_timeout is None:
+                # This is an infinite generator, so it should never exit.
+                yield from backoff_generator
+                return
+
+            # Generate timeout intervals until we have reached the
+            # `etp_timeout`-threshold.
+            csum = 0.0
+            for d in backoff_generator:
+                yield d
+
+                csum += d
+
+                if csum >= etp_timeout:
+                    break
+
+        for ti in timeout_intervals(self.etp_timeout):
+            try:
+                # Wait for an event for `ti` seconds.
+                async with timeout(ti):
+                    await self._recv_events[correlation_id].wait()
+            except TimeoutError:
+                # Check if the receiver task is still running.
+                if self.__recvtask.done():
+                    # Raise any errors by waiting for the task to finish.
+                    await self.__recvtask
+
+                    logger.error(
+                        "Receiver task terminated without errors. This should not happen"
+                    )
+
+                    raise ReceiveWorkerExited
+            else:
+                # Break out of for-loop, and start processing message.
+                break
+        else:
+            # The for-loop finished without breaking. In other words, we have
+            # timed out.
+            assert self.etp_timeout is not None
+            raise TimeoutError(
+                f"Receiver task did not set event within {self.etp_timeout} seconds"
+            )
 
-        # cleanup
-        bodies = self._clear_msg_on_buffer(correlation_id)
+        # Remove event from list of events
+        del self._recv_events[correlation_id]
+        # Read message bodies from buffer.
+        bodies = self._recv_buffer.pop(correlation_id)
 
-        # error handling
+        # Check if there are errors in the received messages.
         errors = self._parse_error_info(bodies)
 
+        # Raise errors in case there are any.
         if len(errors) == 1:
             raise ETPError.from_proto(errors.pop())
         elif len(errors) > 1:
@@ -285,9 +372,14 @@ class ETPClient(ETPConnection):
                 errors.extend(body.errors.values())
         return errors
 
+    async def __aexit__(self, *exc_details) -> None:
+        await self.close(reason="Client exiting")
+
     async def close(self, reason=""):
+        close_session_sent = False
         try:
             await self._send(CloseSession(reason=reason))
+            close_session_sent = True
         except websockets.ConnectionClosed:
             logger.error(
                 "Websockets connection is closed, unable to send a CloseSession-message"
@@ -324,12 +416,23 @@ class ETPClient(ETPConnection):
         # Reading them will speed up the closing of the connection.
         counter = 0
         try:
-            async for msg in self.ws:
-                counter += 1
+            # In some cases the server does not drop the connection after we
+            # have sent the `CloseSession`-message. We therefore add a timeout
+            # to the reading of possibly lost messages.
+            async with timeout(self.etp_timeout or 10):
+                async for msg in self.ws:
+                    counter += 1
         except websockets.ConnectionClosed:
-            # The websockets connection had already closed. Either successfully
+            # The websockets connection has already closed. Either successfully
             # or with an error, but we ignore both cases.
             pass
+        except TimeoutError:
+            if close_session_sent:
+                logger.error(
+                    "Websockets connection was not closed within "
+                    f"{self.etp_timeout or 10} seconds after the "
+                    "`CloseSession`-message was sent"
+                )
 
         if counter > 0:
             logger.error(
@@ -339,26 +442,16 @@ class ETPClient(ETPConnection):
 
         logger.debug("Client closed")
 
-    #
-    #
-    #
-
-    def _clear_msg_on_buffer(self, correlation_id: int):
-        del self._recv_events[correlation_id]
-        return self._recv_buffer.pop(correlation_id)
-
-    def _add_msg_to_buffer(self, msg: Message):
-        self._recv_buffer[msg.header.correlation_id].append(msg.body)
-
-        # NOTE: should we add task to autoclear buffer message if never waited on ?
-        if msg.is_final_msg():
-            # set response on send event
-            self._recv_events[msg.header.correlation_id].set()
-
-    async def __recv__(self):
-        logger.debug("starting recv loop")
+    async def __recv(self):
+        logger.debug("Starting receiver loop")
 
         while True:
+            # We use this way of receiving messages, instead of the `async
+            # for`-pattern, in order to raise all
+            # `websockets.exceptions.ConnectionClosed`-errors (including the
+            # `websockets.exceptions.ConnectionClosedOK` error). In the `async
+            # for`-case a closing code of `1000` (normal closing) just exits
+            # the loop.
             msg_data = await self.ws.recv()
             msg = Message.decode_binary_message(
                 T.cast(bytes, msg_data), ETPClient.generic_transition_table
@@ -369,12 +462,19 @@ class ETPClient(ETPConnection):
                 continue
 
             logger.debug(f"recv {msg.body.__class__.__name__} {repr(msg.header)}")
-            self._add_msg_to_buffer(msg)
+            self._recv_buffer[msg.header.correlation_id].append(msg.body)
+
+            if msg.is_final_msg():
+                # set response on send event
+                self._recv_events[msg.header.correlation_id].set()
 
     #
     # session related
     #
 
+    async def __aenter__(self) -> Self:
+        return await self.request_session()
+
     async def request_session(self):
         # Handshake protocol
         etp_version = Version(major=1, minor=2, revision=0, patch=0)
@@ -390,8 +490,8 @@ class ETPClient(ETPConnection):
 
         msg = await self.send(
             RequestSession(
-                applicationName=SETTINGS.application_name,
-                applicationVersion=SETTINGS.application_version,
+                applicationName=self.application_name,
+                applicationVersion=self.application_version,
                 clientInstanceId=uuid.uuid4(),  # type: ignore
                 requestedProtocols=[
                     SupportedProtocol(
@@ -435,16 +535,23 @@ class ETPClient(ETPConnection):
 
         return msg
 
-    #
+    @staticmethod
+    def assert_response(response: ETPModel, expected_type: T.Type[ETPModel]) -> None:
+        assert isinstance(response, expected_type), (
+            f"Expected {expected_type}, got {type(response)} with content {response}"
+        )
 
     @property
     def max_size(self):
-        return SETTINGS.MaxWebSocketMessagePayloadSize
-        # return self.client_info.getCapability("MaxWebSocketMessagePayloadSize")
+        return self.client_info.getCapability("MaxWebSocketMessagePayloadSize")
 
     @property
     def max_array_size(self):
-        return self.max_size - 512  # maxsize - 512 bytes for header and body
+        if self.max_size <= 3000:
+            raise AttributeError(
+                "The maximum size of a websocket message must be greater than 3000"
+            )
+        return self.max_size - 3000  # maxsize - 3000 bytes for header and body
 
     @property
     def timestamp(self):
@@ -455,12 +562,12 @@ class ETPClient(ETPConnection):
             raise Exception("Max one / in dataspace name")
         return DataspaceURI.from_name(ds)
 
-    def list_objects(self, dataspace_uri: DataspaceURI, depth: int = 1) -> list:
+    def list_objects(self, dataspace_uri: DataspaceURI | str, depth: int = 1) -> list:
         return self.send(
             GetResources(
                 scope=ContextScopeKind.TARGETS_OR_SELF,
                 context=ContextInfo(
-                    uri=dataspace_uri.raw_uri,
+                    uri=str(dataspace_uri),
                     depth=depth,
                     dataObjectTypes=[],
                     navigableEdges=RelationshipKind.PRIMARY,
@@ -549,33 +656,67 @@ class ETPClient(ETPConnection):
         )
         return response.success
 
-    #
-    # data objects
-    #
-
     async def get_data_objects(self, *uris: T.Union[DataObjectURI, str]):
-        _uris = list(map(str, uris))
+        tasks = []
+        for uri in uris:
+            task = self.send(GetDataObjects(uris={str(uri): str(uri)}))
+            tasks.append(task)
 
-        msg = await self.send(GetDataObjects(uris=dict(zip(_uris, _uris))))
-        assert isinstance(msg, GetDataObjectsResponse), "Expected dataobjectsresponse"
-        assert len(msg.data_objects) == len(_uris), (
-            "Here we assume that all three objects fit in a single record"
-        )
+        responses = await asyncio.gather(*tasks)
+        assert len(responses) == len(uris)
+
+        data_objects = []
+        errors = []
+        for uri, response in zip(uris, responses):
+            if not isinstance(response, GetDataObjectsResponse):
+                errors.append(
+                    TypeError(
+                        "Expected GetDataObjectsResponse, got "
+                        f"{response.__class__.__name} with content: {response}",
+                    )
+                )
+            data_objects.append(response.data_objects[str(uri)])
+
+        if len(errors) > 0:
+            raise ExceptionGroup(
+                f"There were {len(errors)} errors in ETPClient.get_data_objects",
+                errors,
+            )
 
-        return [msg.data_objects[u] for u in _uris]
+        return data_objects
 
     async def put_data_objects(self, *objs: DataObject):
-        response = await self.send(
-            PutDataObjects(
-                data_objects={f"{p.resource.name} - {p.resource.uri}": p for p in objs},
+        tasks = []
+        for obj in objs:
+            task = self.send(
+                PutDataObjects(
+                    data_objects={f"{obj.resource.name} - {obj.resource.uri}": obj},
+                ),
             )
-        )
+            tasks.append(task)
 
-        assert isinstance(response, PutDataObjectsResponse), (
-            "Expected PutDataObjectsResponse"
-        )
+        responses = await asyncio.gather(*tasks)
 
-        return response.success
+        errors = []
+        for response in responses:
+            if not isinstance(response, PutDataObjectsResponse):
+                errors.append(
+                    TypeError(
+                        "Expected PutDataObjectsResponse, got "
+                        f"{response.__class__.__name} with content: {response}",
+                    )
+                )
+        if len(errors) > 0:
+            raise ExceptionGroup(
+                f"There were {len(errors)} errors in ETPClient.put_data_objects",
+                errors,
+            )
+
+        sucesses = {}
+        for response in responses:
+            sucesses = {**sucesses, **response.success}
+
+        return sucesses
 
     async def get_resqml_objects(
         self, *uris: T.Union[DataObjectURI, str]
@@ -766,7 +907,12 @@ class ETPClient(ETPConnection):
     async def get_array(self, uid: DataArrayIdentifier):
         # Check if we can download the full array in one go.
         (meta,) = await self.get_array_metadata(uid)
-        if utils_arrays.get_transport_array_size(meta) > self.max_array_size:
+        if (
+            utils_arrays.get_transport_array_size(
+                meta.transport_array_type, meta.dimensions
+            )
+            > self.max_array_size
+        ):
             return await self._get_array_chunked(uid)
 
         response = await self.send(
@@ -779,6 +925,250 @@ class ETPClient(ETPConnection):
         arrays = list(response.data_arrays.values())
         return utils_arrays.get_numpy_array_from_etp_data_array(arrays[0])
 
+    async def download_array(
+        self,
+        epc_uri: str | DataObjectURI,
+        path_in_resource: str,
+    ) -> npt.NDArray[utils_arrays.LogicalArrayDTypes]:
+        # Create identifier for the data.
+        dai = DataArrayIdentifier(
+            uri=str(epc_uri),
+            path_in_resource=path_in_resource,
+        )
+
+        response = await self.send(
+            GetDataArrayMetadata(data_arrays={dai.path_in_resource: dai}),
+        )
+
+        self.assert_response(response, GetDataArrayMetadataResponse)
+        assert (
+            len(response.array_metadata) == 1
+            and dai.path_in_resource in response.array_metadata
+        )
+
+        metadata = response.array_metadata[dai.path_in_resource]
+
+        # Check if we can download the full array in a single message.
+        if (
+            utils_arrays.get_transport_array_size(
+                metadata.transport_array_type, metadata.dimensions
+            )
+            >= self.max_array_size
+        ):
+            transport_dtype = utils_arrays.get_dtype_from_any_array_type(
+                metadata.transport_array_type,
+            )
+            # NOTE: The logical array type is not yet supported by the
+            # open-etp-server. As such the transport array type will be actual
+            # array type used. We only add this call to prepare for when it
+            # will be used.
+            logical_dtype = utils_arrays.get_dtype_from_any_logical_array_type(
+                metadata.logical_array_type,
+            )
+            if logical_dtype != np.dtype(np.bool_):
+                # If this debug message is triggered we should test the
+                # mapping.
+                logger.debug(
+                    "Logical array type has changed: "
+                    f"{metadata.logical_array_type = }, with {logical_dtype = }"
+                )
+
+            # Create a buffer for the data.
+            data = np.zeros(metadata.dimensions, dtype=transport_dtype)
+
+            # Get list with starting indices in each block, and a list with the
+            # number of elements along each axis for each block.
+            block_starts, block_counts = utils_arrays.get_array_block_sizes(
+                data.shape, data.dtype, self.max_array_size
+            )
+
+            def data_subarrays_key(pir: str, i: int) -> str:
+                return pir + f" ({i})"
+
+            tasks = []
+            for i, (starts, counts) in enumerate(zip(block_starts, block_counts)):
+                task = self.send(
+                    GetDataSubarrays(
+                        data_subarrays={
+                            data_subarrays_key(
+                                dai.path_in_resource, i
+                            ): GetDataSubarraysType(
+                                uid=dai,
+                                starts=starts,
+                                counts=counts,
+                            ),
+                        },
+                    ),
+                )
+                tasks.append(task)
+
+            responses = await asyncio.gather(*tasks)
+
+            data_blocks = []
+            for i, response in enumerate(responses):
+                self.assert_response(response, GetDataSubarraysResponse)
+                assert (
+                    len(response.data_subarrays) == 1
+                    and data_subarrays_key(dai.path_in_resource, i)
+                    in response.data_subarrays
+                )
+
+                data_block = utils_arrays.get_numpy_array_from_etp_data_array(
+                    response.data_subarrays[
+                        data_subarrays_key(dai.path_in_resource, i)
+                    ],
+                )
+                data_blocks.append(data_block)
+
+            for data_block, starts, counts in zip(
+                data_blocks, block_starts, block_counts
+            ):
+                # Create slice-objects for each block.
+                slices = tuple(
+                    map(
+                        lambda s, c: slice(s, s + c),
+                        np.array(starts).astype(int),
+                        np.array(counts).astype(int),
+                    )
+                )
+                data[slices] = data_block
+
+            # Return after fetching all sub arrays.
+            return data
+
+        # Download the full array in one go.
+        response = await self.send(
+            GetDataArrays(data_arrays={dai.path_in_resource: dai}),
+        )
+
+        self.assert_response(response, GetDataArraysResponse)
+        assert (
+            len(response.data_arrays) == 1
+            and dai.path_in_resource in response.data_arrays
+        )
+
+        return utils_arrays.get_numpy_array_from_etp_data_array(
+            response.data_arrays[dai.path_in_resource]
+        )
+
+    async def upload_array(
+        self,
+        epc_uri: str | DataObjectURI,
+        path_in_resource: str,
+        data: npt.NDArray[utils_arrays.LogicalArrayDTypes],
+    ) -> None:
+        # Fetch ETP logical and transport array types
+        logical_array_type, transport_array_type = (
+            utils_arrays.get_logical_and_transport_array_types(data.dtype)
+        )
+
+        # Create identifier for the data.
+        dai = DataArrayIdentifier(
+            uri=str(epc_uri),
+            path_in_resource=path_in_resource,
+        )
+
+        # Get current time as a UTC-timestamp.
+        now = self.timestamp
+
+        # Allocate space on server for the array.
+        response = await self.send(
+            PutUninitializedDataArrays(
+                data_arrays={
+                    dai.path_in_resource: PutUninitializedDataArrayType(
+                        uid=dai,
+                        metadata=DataArrayMetadata(
+                            dimensions=list(data.shape),
+                            transport_array_type=transport_array_type,
+                            logical_array_type=logical_array_type,
+                            store_last_write=now,
+                            store_created=now,
+                        ),
+                    ),
+                },
+            ),
+        )
+
+        self.assert_response(response, PutUninitializedDataArraysResponse)
+        assert len(response.success) == 1 and dai.path_in_resource in response.success
+
+        # Check if we can upload the entire array in go, or if we need to
+        # upload it in smaller blocks.
+        if data.nbytes > self.max_array_size:
+            tasks = []
+
+            # Get list with starting indices in each block, and a list with the
+            # number of elements along each axis for each block.
+            block_starts, block_counts = utils_arrays.get_array_block_sizes(
+                data.shape, data.dtype, self.max_array_size
+            )
+
+            for starts, counts in zip(block_starts, block_counts):
+                # Create slice-objects for each block.
+                slices = tuple(
+                    map(
+                        lambda s, c: slice(s, s + c),
+                        np.array(starts).astype(int),
+                        np.array(counts).astype(int),
+                    )
+                )
+
+                # Slice the array, and convert to the relevant ETP-array type.
+                # Note in the particular the extra `.data`-after the call. The
+                # data should not be of type `DataArray`, but `AnyArray`, so we
+                # need to fetch it from the `DataArray`.
+                etp_subarray_data = utils_arrays.get_etp_data_array_from_numpy(
+                    data[slices]
+                ).data
+
+                # Create an asynchronous task to upload a block to the
+                # ETP-server.
+                task = self.send(
+                    PutDataSubarrays(
+                        data_subarrays={
+                            dai.path_in_resource: PutDataSubarraysType(
+                                uid=dai,
+                                data=etp_subarray_data,
+                                starts=starts,
+                                counts=counts,
+                            ),
+                        },
+                    ),
+                )
+                tasks.append(task)
+
+            # Upload all blocks.
+            responses = await asyncio.gather(*tasks)
+
+            # Check for successful responses.
+            for response in responses:
+                self.assert_response(response, PutDataSubarraysResponse)
+                assert (
+                    len(response.success) == 1
+                    and dai.path_in_resource in response.success
+                )
+
+            # Return after uploading all sub arrays.
+            return
+
+        # Convert NumPy data-array to an ETP-transport array.
+        etp_array_data = utils_arrays.get_etp_data_array_from_numpy(data)
+
+        # Pass entire array in one message.
+        response = await self.send(
+            PutDataArrays(
+                data_arrays={
+                    dai.path_in_resource: PutDataArraysType(
+                        uid=dai,
+                        array=etp_array_data,
+                    ),
+                }
+            )
+        )
+
+        self.assert_response(response, PutDataArraysResponse)
+        assert len(response.success) == 1 and dai.path_in_resource in response.success
+
     async def put_array(
         self,
         uid: DataArrayIdentifier,
@@ -858,7 +1248,7 @@ class ETPClient(ETPConnection):
         counts = np.array(counts).astype(np.int64)  # len = 2
         ends = starts + counts  # len = 2
 
-        slices = tuple(map(lambda se: slice(se[0], se[1]), zip(starts, ends)))
+        slices = tuple(map(lambda s, e: slice(s, e), starts, ends))
         dataarray = utils_arrays.get_etp_data_array_from_numpy(data[slices])
         payload = PutDataSubarraysType(
             uid=uid,
@@ -881,17 +1271,19 @@ class ETPClient(ETPConnection):
         assert len(response.success) == 1, "expected one success"
         return response.success
 
-    #
-    # chunked get array - ETP will not chunk response - so we need to do it manually
-    #
-
     def _get_chunk_sizes(
         self, shape, dtype: np.dtype[T.Any] = np.dtype(np.float32), offset=0
     ):
+        warnings.warn(
+            "This function is deprecated and will be removed in a later version of "
+            "pyetp. The replacement is located via the import "
+            "`from pyetp.utils_arrays import get_array_block_sizes`.",
+            DeprecationWarning,
+            stacklevel=2,
+        )
         shape = np.array(shape)
 
         # capsize blocksize
-        # remove 512 bytes for headers and body
         max_items = self.max_array_size / dtype.itemsize
         block_size = np.power(max_items, 1.0 / len(shape))
         block_size = min(2048, int(block_size // 2) * 2)
@@ -1037,7 +1429,13 @@ class connect:
             open_timeout=None,
         )
 
-        self.client = ETPClient(self.ws, timeout=self.timeout)
+        self.client = ETPClient(
+            self.ws,
+            etp_timeout=self.timeout,
+            max_message_size=SETTINGS.MaxWebSocketMessagePayloadSize,
+            application_name=SETTINGS.application_name,
+            application_version=SETTINGS.application_version,
+        )
 
         try:
             await self.client.request_session()
@@ -1052,3 +1450,72 @@ class connect:
     async def __aexit__(self, exc_type, exc: Exception, tb: TracebackType):
         await self.client.close()
         await self.ws.close()
+
+
+@contextlib.asynccontextmanager
+async def etp_connect(
+    uri: str,
+    data_partition_id: str | None = None,
+    authorization: str | None = None,
+    etp_timeout: float = 10.0,
+    max_message_size: float = 2**20,
+) -> ETPClient:
+    additional_headers = {}
+
+    if authorization is not None:
+        additional_headers["Authorization"] = authorization
+    if data_partition_id is not None:
+        additional_headers["data-partition-id"] = data_partition_id
+
+    subprotocols = ["etp12.energistics.org"]
+
+    async with (
+        websockets.connect(
+            uri=uri,
+            subprotocols=subprotocols,
+            max_size=max_message_size,
+            additional_headers=additional_headers,
+        ) as ws,
+        ETPClient(
+            ws=ws,
+            etp_timeout=etp_timeout,
+            max_message_size=max_message_size,
+        ) as etp_client,
+    ):
+        yield etp_client
+
+
+async def etp_persistent_connect(
+    uri: str,
+    data_partition_id: str | None = None,
+    authorization: str | None = None,
+    etp_timeout: float = 10.0,
+    max_message_size: float = 2**20,
+) -> AsyncGenerator[ETPClient]:
+    additional_headers = {}
+
+    if authorization is not None:
+        additional_headers["Authorization"] = authorization
+    if data_partition_id is not None:
+        additional_headers["data-partition-id"] = data_partition_id
+
+    subprotocols = ["etp12.energistics.org"]
+    async for ws in websockets.connect(
+        uri=uri,
+        subprotocols=subprotocols,
+        max_size=max_message_size,
+        additional_headers=additional_headers,
+    ):
+        try:
+            async with ETPClient(
+                ws=ws,
+                etp_timeout=etp_timeout,
+                max_message_size=max_message_size,
+            ) as etp_client:
+                yield etp_client
+        except websockets.ConnectionClosed as e:
+            logger.info(
+                f"Websockets connection closed with message '{e}'. Starting new "
+                "connection"
+            )
+            continue

pyetp/utils_arrays.py

@@ -1,3 +1,4 @@
+import sys
 import typing as T
 
 import numpy as np
@@ -15,9 +16,6 @@ from etptypes.energistics.etp.v12.datatypes.array_of_long import ArrayOfLong
 from etptypes.energistics.etp.v12.datatypes.data_array_types.data_array import (
     DataArray,
 )
-from etptypes.energistics.etp.v12.datatypes.data_array_types.data_array_metadata import (
-    DataArrayMetadata,
-)
 
 SUPPORTED_ARRAY_TYPES: T.TypeAlias = (
     ArrayOfFloat | ArrayOfBoolean | ArrayOfInt | ArrayOfLong | ArrayOfDouble
@@ -51,6 +49,14 @@ _ANY_LOGICAL_ARRAY_TYPE_MAP: dict[npt.DTypeLike, AnyLogicalArrayType] = {
     np.dtype(">f8"): AnyLogicalArrayType.ARRAY_OF_DOUBLE64_BE,
 }
 
+valid_logical_array_dtypes = list(_ANY_LOGICAL_ARRAY_TYPE_MAP)
+if (sys.version_info.major, sys.version_info.minor) == (3, 10):
+    LogicalArrayDTypes: T.TypeAlias = T.Union[
+        tuple(v.type for v in valid_logical_array_dtypes)
+    ]
+else:
+    LogicalArrayDTypes: T.TypeAlias = T.Union[tuple(valid_logical_array_dtypes)]
+
 _INV_ANY_LOGICAL_ARRAY_TYPE_MAP: dict[AnyLogicalArrayType, npt.DTypeLike] = {
     v: k for k, v in _ANY_LOGICAL_ARRAY_TYPE_MAP.items()
 }
@@ -91,7 +97,7 @@ _ANY_ARRAY_TYPE_MAP: dict[npt.DTypeLike, AnyArrayType] = {
     np.dtype("<f4"): AnyArrayType.ARRAY_OF_FLOAT,
     np.dtype("<f8"): AnyArrayType.ARRAY_OF_DOUBLE,
 }
-valid_dtypes = list(_ANY_ARRAY_TYPE_MAP)
+valid_any_array_dtypes = list(_ANY_ARRAY_TYPE_MAP)
 
 _INV_ANY_ARRAY_TYPE_MAP: dict[AnyArrayType, npt.DTypeLike] = {
     AnyArrayType.ARRAY_OF_BOOLEAN: np.dtype(np.bool_),
@@ -121,7 +127,7 @@ _INV_ANY_ARRAY_MAP: dict[SUPPORTED_ARRAY_TYPES, AnyArrayType] = {
 
 
 def check_if_array_is_valid_dtype(array: npt.NDArray[T.Any]) -> bool:
-    return array.dtype in valid_dtypes
+    return array.dtype in valid_any_array_dtypes
 
 
 def get_valid_dtype_cast(array: npt.NDArray[T.Any]) -> npt.DTypeLike:
@@ -210,9 +216,11 @@ def get_etp_data_array_from_numpy(data: npt.NDArray) -> DataArray:
     )
 
 
-def get_transport_array_size(metadata: DataArrayMetadata) -> int:
-    dtype = _INV_ANY_ARRAY_TYPE_MAP[metadata.transport_array_type]
-    return int(np.prod(metadata.dimensions) * dtype.itemsize)
+def get_transport_array_size(
+    transport_array_type: AnyArrayType, dimensions: list[int] | tuple[int]
+) -> int:
+    dtype = _INV_ANY_ARRAY_TYPE_MAP[transport_array_type]
+    return int(np.prod(dimensions) * dtype.itemsize)
 
 
 def get_dtype_from_any_array_class(cls: AnyArray) -> npt.DTypeLike:
@@ -235,11 +243,18 @@ def get_dtype_from_any_array_class(cls: AnyArray) -> npt.DTypeLike:
     raise TypeError(f"Class {cls} is not a valid array class")
 
 
-def get_dtype_from_any_array_type(_type: T.Union[AnyArrayType | str]) -> npt.DTypeLike:
+def get_dtype_from_any_array_type(_type: AnyArrayType | str) -> npt.DTypeLike:
     enum_name = AnyArrayType(_type)
     return _INV_ANY_ARRAY_TYPE_MAP[enum_name]
 
 
+def get_dtype_from_any_logical_array_type(
+    _type: AnyLogicalArrayType | str,
+) -> npt.DTypeLike:
+    enum_name = AnyLogicalArrayType(_type)
+    return _INV_ANY_LOGICAL_ARRAY_TYPE_MAP[enum_name]
+
+
 def get_numpy_array_from_etp_data_array(
     data_array: DataArray,
 ) -> npt.NDArray[
@@ -258,3 +273,67 @@ def get_numpy_array_from_etp_data_array(
     return np.array(data_array.data.item.values, dtype=dtype).reshape(
         data_array.dimensions
     )
+
+
+def get_array_block_sizes(
+    shape: tuple[int], dtype: npt.DTypeLike, max_array_size: int
+) -> tuple[list[list[int]], list[list[int]]]:
+    # Total size of array in bytes.
+    array_size = int(np.prod(shape) * dtype.itemsize)
+    # Calculate the minimum number of blocks needed (if the array was flat).
+    num_blocks = int(np.ceil(array_size / max_array_size))
+
+    # Check if we can split on the first axis.
+    if num_blocks > shape[0]:
+        assert len(shape) > 1
+        # Recursively get block sizes on higher axes.
+        starts, counts = get_array_block_sizes(shape[1:], dtype, max_array_size)
+        # Repeat starts and counts from higher axes for each axis 0.
+        starts = [[i] + s for i in range(shape[0]) for s in starts]
+        counts = [[1] + c for i in range(shape[0]) for c in counts]
+
+        return starts, counts
+
+    # Count the number of axis elements (e.g., rows for a 2d-array) in each
+    # block, and count the number of blocks that remain.
+    num_elements_in_block, num_remainder = divmod(shape[0], num_blocks)
+
+    # Get the number of extra blocks needed to fill in the remaining elements.
+    num_extra_blocks = num_remainder // num_elements_in_block + int(
+        num_remainder % num_elements_in_block > 0
+    )
+    # Count the number of elements in the last block.
+    num_elements_in_last_block = num_remainder % num_elements_in_block
+    # Increase the number of blocks to fit the remaining elements.
+    num_blocks += num_extra_blocks
+
+    # Verify that we still have more axis elements than blocks.
+    assert num_blocks <= shape[0]
+
+    # Set up the number of axis elements in each block.
+    axis_counts = np.ones(num_blocks, dtype=int) * num_elements_in_block
+    if num_elements_in_last_block > 0:
+        assert num_elements_in_last_block < num_elements_in_block
+        # Alter the last block with the remaining number of elements.
+        axis_counts[-1] = num_elements_in_last_block
+
+    # Create an array with starting indices for each block and a corresponding
+    # array with the number of elements in each block.
+    starts = np.zeros((num_blocks, len(shape)), dtype=int)
+    counts = np.zeros_like(starts)
+
+    # Sum up the number of element counts to get the starting index in each
+    # block (starting at 0).
+    starts[1:, 0] = np.cumsum(axis_counts[:-1])
+
+    # The axis_counts already lists the number of elements in the first axis,
+    # so we only add the length of each remaining axis as the counts for the
+    # last axes.
+    counts[:, 0] = axis_counts
+    counts[:, 1:] = shape[1:]
+
+    # Check that no block exceeds the maximum size.
+    count_size = np.prod(counts, axis=1) * dtype.itemsize
+    assert np.all(count_size - max_array_size <= 0)
+
+    return starts.tolist(), counts.tolist()

{pyetp-0.0.43.dist-info → pyetp-0.0.44.dist-info}/METADATA

@@ -1,20 +1,21 @@
 Metadata-Version: 2.4
 Name: pyetp
-Version: 0.0.43
+Version: 0.0.44
 Summary: Interface with OSDU RDDMS using ETP protocol
 Author-email: Adam Cheng <52572642+adamchengtkc@users.noreply.github.com>
 License-Expression: Apache-2.0
 Project-URL: homepage, https://github.com/equinor/pyetp
 Classifier: Development Status :: 3 - Alpha
+Requires-Python: >=3.10
 Description-Content-Type: text/markdown
 License-File: LICENSE.md
-Requires-Dist: numpy>=1.26.3
-Requires-Dist: websockets>=12.0
+Requires-Dist: numpy>=2.0
+Requires-Dist: websockets>=15.0
 Requires-Dist: lxml>=4.9.4
 Requires-Dist: pydantic>=1.10
 Requires-Dist: async-timeout>=5.0
 Requires-Dist: xtgeo>=4.0.0
-Requires-Dist: xsdata>=24.3.1
+Requires-Dist: xsdata>=25.4
 Requires-Dist: etpproto>=1.0.7
 Dynamic: license-file
 
@@ -29,6 +30,7 @@ This package is published to PyPI, and can be installed via:
 ```bash
 pip install pyetp
 ```
+The library is tested against Python versions 3.10, 3.11, 3.12 and 3.13.
 
 ## Local development
 Locally we suggest setting up a virtual environment, and installing the latest

{pyetp-0.0.43.dist-info → pyetp-0.0.44.dist-info}/RECORD

@@ -1,13 +1,13 @@
-pyetp/__init__.py,sha256=Vu3_qz0AazlD4Q6ZLGdQVqZ7lhzqNLivpQRS3owD8SY,251
-pyetp/_version.py,sha256=qJphZkKjg5qeGMzBrtqKwyJZQtcx3oaEh-0t6ejMuVo,706
-pyetp/client.py,sha256=kZsPi2OPvbQKX6IVoAsrdZ9_ZDbZYEhHuHOhfbf3XxQ,37427
+pyetp/__init__.py,sha256=_mJRvg6XJgzHeRWd_AsBNonvWUBD3jroZ9MLehPRHP4,337
+pyetp/_version.py,sha256=Zrt00MLeXbWP2ZsKSWktjZjWuHHC_-KLVmAWLgE0o-g,706
+pyetp/client.py,sha256=BF-_sm9j_SbJsQZ7tK8kp5Ekl7FsNtMHAZFJAwHuTMQ,54475
 pyetp/config.py,sha256=uGEx6n-YF7Rtgwckf0ovRKNOKgCUsiQx4IA0Tyiqafk,870
 pyetp/resqml_objects.py,sha256=j00e8scSh-yYv4Lpp9WjzLiaKqJWd5Cs7ROcJcxxFVw,50721
 pyetp/types.py,sha256=zOfUzEQcgBvveWJyM6dD7U3xJ4SCkWElewNL0Ml-PPY,6761
 pyetp/uri.py,sha256=Y05vuTO-XOurgDavBeoPGOidALoCKjCBIb7YHmfbAco,3115
-pyetp/utils_arrays.py,sha256=bLrb8H8TMjbpNLo0Zb9xYRUnM2W4BOf7WxOxWK4fcTc,10218
+pyetp/utils_arrays.py,sha256=BJMiwx0mnSUjAOARCfCvWr_rz6N38e7kSWi24I1bCLA,13447
 pyetp/utils_xml.py,sha256=i11Zv2PuxW-ejEXxfGPYHvuJv7EHPYXwvgoe1U8fOUM,6826
-pyetp-0.0.43.dist-info/licenses/LICENSE.md,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
+pyetp-0.0.44.dist-info/licenses/LICENSE.md,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
 resqml_objects/__init__.py,sha256=ilimTrQScFryxHrfZwZURVpW0li19bVnayUcHR7S_Fs,183
 resqml_objects/epc_readers.py,sha256=InYMlwjiZZRG9poQlWKFEOUNJmowGahXoNu_GniOaBw,3458
 resqml_objects/parsers.py,sha256=UZdBi3QmBq4ejwKI9Dse_lMqL5Bx2s3QNTs8TKS4fO0,427
@@ -15,7 +15,7 @@ resqml_objects/serializers.py,sha256=JqRO6D6ExT5DrVyiNwBgNW28108f6spvxiVqJU0D9mc
 resqml_objects/v201/__init__.py,sha256=yL3jWgkyGAzr-wt_WJDV_eARE75NoA6SPEKBrKM4Crk,51630
 resqml_objects/v201/generated.py,sha256=Se0eePS6w25sfmnp2UBSkgzDGJ9c9Y2QqJgDRUTt_-Q,769527
 resqml_objects/v201/utils.py,sha256=WiywauiJRBWhdjUvbKhpltRjoBX3qWd7qQ0_FAmIzUc,1442
-pyetp-0.0.43.dist-info/METADATA,sha256=nFszNnd-08BucrOKHaFZaB-kMiFz_vqxk6QTWdl04bU,2587
-pyetp-0.0.43.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-pyetp-0.0.43.dist-info/top_level.txt,sha256=NrdXbidkT5QR4NjH6nv2Frixknqse3jZq7bnqNdVb5k,21
-pyetp-0.0.43.dist-info/RECORD,,
+pyetp-0.0.44.dist-info/METADATA,sha256=51q_BHQ-1OSNjt021B3R8zveiDViTHWapaaSrgYVy6U,2679
+pyetp-0.0.44.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+pyetp-0.0.44.dist-info/top_level.txt,sha256=NrdXbidkT5QR4NjH6nv2Frixknqse3jZq7bnqNdVb5k,21
+pyetp-0.0.44.dist-info/RECORD,,