Hide navigation sidebar
Hide table of contents sidebar
Skip to content
ASDF Projects Tutorials Community
Back to top
Toggle table of contents sidebar

# Source code for asdf.extension._serialization_context

    import enum

    from asdf._block.key import Key as BlockKey
    from asdf._helpers import validate_version
    from asdf.extension._extension import ExtensionProxy

    class SerializationContext:
        """
        Container for parameters of the current (de)serialization.

        This class should not be instantiated directly and instead
        will be created by the AsdfFile object and provided to extension
        classes (like Converters) via method arguments.
        """

        def __init__(self, version, extension_manager, url, blocks):
            self._version = validate_version(version)
            self._extension_manager = extension_manager
            self._url = url
            self._blocks = blocks
            self._obj = None

            self.__extensions_used = set()

        @property
        def url(self):
            """
            The URL (if any) of the file being read or written.

            Used to compute relative locations of external files referenced by this
            ASDF file. The URL will not exist in some cases (e.g. when the file is
            written to an `io.BytesIO`).

            Returns
            -------
            str or None
            """
            return self._url

        @property
        def version(self):
            """
            Get the ASDF Standard version.

            Returns
            -------
            str
            """
            return self._version

        @property
        def extension_manager(self):
            """
            Get the ExtensionManager for enabled extensions.

            Returns
            -------
            asdf.extension.ExtensionManager
            """
            return self._extension_manager

        def _mark_extension_used(self, extension):
            """
            Note that an extension was used when reading or writing the file.

            Parameters
            ----------
            extension : asdf.extension.Extension
            """
            self.__extensions_used.add(ExtensionProxy.maybe_wrap(extension))

        @property
        def _extensions_used(self):
            """
            Get the set of extensions that were used when reading or writing the file.

            Returns
            -------
            set of asdf.extension.Extension
            """
            return self.__extensions_used

        def get_block_data_callback(self, index, key=None):
            """
            Generate a callable that when called will read data
            from an ASDF block at the provided index.

            Parameters
            ----------
            index : int
                Index of ASDF block.

            key : BlockKey, optional
                BlockKey generated using self.generate_block_key. Only
                needed for a Converter that uses multiple blocks.

            Returns
            -------
            callback : callable
                A callable that when called (with no arguments) returns
                the block data as a one dimensional array of uint8
            """
            raise NotImplementedError("abstract")

        def find_available_block_index(self, data_callback, key=None):
            """
            Find the index of an available ASDF block to write data.

            This is typically used inside asdf.extension.Converter.to_yaml_tree.

            Parameters
            ----------
            data_callback: callable
                Callable that when called will return data (ndarray) that will
                be written to a block.

            key : BlockKey, optional
                BlockKey generated using self.generate_block_key. Only
                needed for a Converter that uses multiple blocks.

            Returns
            -------
            block_index: int
                Index of the ASDF block where data returned from
                data_callback will be written.
            """
            raise NotImplementedError("abstract")

        def generate_block_key(self):
            """
            Generate a BlockKey used for Converters that wish to use
            multiple blocks

            Returns
            -------
            key : BlockKey
                A hashable object that will be associated with the
                serialized/deserialized object and can be used to
                access multiple blocks within a Converter
            """
            raise NotImplementedError("abstract")

        def assign_object(self, obj):
            self._obj = obj

        def assign_blocks(self):
            pass

        def set_array_storage(self, arr, array_storage):
            """
            Set the block type to use for the given array data.

            Parameters
            ----------
            arr : numpy.ndarray
                The array to set.  If multiple views of the array are in
                the tree, only the most recent block type setting will be
                used, since all views share a single block.

            array_storage : str
                Must be one of:

                - ``internal``: The default.  The array data will be
                  stored in a binary block in the same ASDF file.

                - ``external``: Store the data in a binary block in a
                  separate ASDF file.

                - ``inline``: Store the data as YAML inline in the tree.
            """
            self._blocks._set_array_storage(arr, array_storage)

        def get_array_storage(self, arr):
            """
            Get the block type for the given array data.

            Parameters
            ----------
            arr : numpy.ndarray
            """
            return self._blocks._get_array_storage(arr)

        def set_array_compression(self, arr, compression, **compression_kwargs):
            """
            Set the compression to use for the given array data.

            Parameters
            ----------
            arr : numpy.ndarray
                The array to set.  If multiple views of the array are in
                the tree, only the most recent compression setting will be
                used, since all views share a single block.

            compression : str or None
                Must be one of:

                - ``''`` or `None`: no compression

                - ``zlib``: Use zlib compression

                - ``bzp2``: Use bzip2 compression

                - ``lz4``: Use lz4 compression

                - ``input``: Use the same compression as in the file read.
                  If there is no prior file, acts as None.

            """
            self._blocks._set_array_compression(arr, compression, **compression_kwargs)

        def get_array_compression(self, arr):
            """
            Get the compression type for the given array data.

            Parameters
            ----------
            arr : numpy.ndarray

            Returns
            -------
            compression : str or None
            """
            return self._blocks._get_array_compression(arr)

        def get_array_compression_kwargs(self, arr):
            """ """
            return self._blocks._get_array_compression_kwargs(arr)

        def set_array_save_base(self, arr, save_base):
            """
            Set the ``save_base`` option for ``arr``. When ``arr`` is
            written to a file, if ``save_base`` is ``True`` the base array
            for ``arr`` will be saved.

            Note that similar to other array options this setting is linked
            to the base array if ``arr`` is a view.

            Parameters
            ----------
            arr : numpy.ndarray

            save_base : bool or None
                if ``None`` the ``default_array_save_base`` value from asdf
                config will be used
            """
            self._blocks._set_array_save_base(arr, save_base)

        def get_array_save_base(self, arr):
            """
            Returns the ``save_base`` option for ``arr``. When ``arr`` is
            written to a file, if ``save_base`` is ``True`` the base array
            for ``arr`` will be saved.

            Parameters
            ----------
            arr : numpy.ndarray

            Returns
            -------
            save_base : bool
            """
            return self._blocks._get_array_save_base(arr)

    class ReadBlocksContext(SerializationContext):
        """
        Perform deserialization (reading) with a `SerializationContext`.

        To allow for block access, `ReadBlocksContext` implements:
            - `SerializationContext.generate_block_key`
            - `SerializationContext.get_block_data_callback`
        and tracks which blocks (and keys) are accessed, assigning them
        to the deserialized object after `assign_object` and
        `assign_blocks` are called.
        """

        def __init__(self, *args, **kwargs):
            super().__init__(*args, **kwargs)
            self.assign_object(None)

        def assign_object(self, obj):
            super().assign_object(obj)
            if obj is None:
                self._cb = None
                self._keys_to_assign = {}

        def assign_blocks(self):
            super().assign_blocks()
            if self._cb is not None:
                self._blocks._data_callbacks.assign_object(self._obj, self._cb)
            for key, cb in self._keys_to_assign.items():
                if cb is None:
                    msg = "Converter generated a key that was never used"
                    raise OSError(msg)
                # now that we have an object, make the key valid
                key._assign_object(self._obj)

                # assign the key to the callback
                self._blocks._data_callbacks.assign_object(key, cb)
            # now that we've assigned blocks, remove the reference to the
            # assigned object
            self.assign_object(None)

        def get_block_data_callback(self, index, key=None):
            if key is None:
                if self._cb is not None:
                    # this operation has already accessed a block without using
                    # a key so check if the same index was accessed
                    if self._cb._index == index:
                        return self._cb
                    msg = "Converters accessing >1 block must provide a key for each block"
                    raise OSError(msg)
                self._cb = self._blocks._get_data_callback(index)
                return self._cb

            if self._keys_to_assign.get(key, None) is not None:
                return self._keys_to_assign[key]

            cb = self._blocks._get_data_callback(index)
            # mark this as a key to later assign
            self._keys_to_assign[key] = cb
            return cb

        def generate_block_key(self):
            key = BlockKey()
            self._keys_to_assign[key] = None
            return key

    class WriteBlocksContext(SerializationContext):
        """
        Perform serialization (writing) with a `SerializationContext`.

        To allow for block access, `WriteBlocksContext` implements:
            - `SerializationContext.generate_block_key`
            - `SerializationContext.find_available_block_index`
        and assigns any accessed blocks (and keys) to the object
        being serialized.
        """

        def find_available_block_index(self, data_callback, key=None):
            if key is None:
                key = self._obj
            return self._blocks.make_write_block(data_callback, None, key)

        def generate_block_key(self):
            return BlockKey(self._obj)

    class BlockAccess(enum.Enum):
        """
        Block access enumerated values that define
        how a SerializationContext can access ASDF blocks.
        """

        NONE = SerializationContext
        WRITE = WriteBlocksContext
        READ = ReadBlocksContext

    def create(asdf_file, block_access=BlockAccess.NONE):
        """
        Create a SerializationContext instance (or subclass) using
        an AsdfFile instance, asdf_file.

        Parameters
        ----------
        asdf_file : asdf.AsdfFile

        block_access : BlockAccess, optional
            Defaults to BlockAccess.NONE
        """
        return block_access.value(asdf_file.version_string, asdf_file.extension_manager, asdf_file.uri, asdf_file._blocks)