Source code for yastn._split_combine_dict

# Copyright 2025 The YASTN Authors. All Rights Reserved.
#
# 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.
# ==============================================================================
""" yastn.split_data_and_meta() and yastn.combine_data_and_meta() """
from __future__ import annotations

__all__ = ['split_data_and_meta', 'combine_data_and_meta']




def split_data_and_meta(d: dict, squeeze=False) -> tuple[tuple['numpy.array' | 'torch.tensor'], dict]:
    """
    Split a dictionary generated by `to_dict` methods into a tuple containing data array,
    and a dictionary `meta` with remaining information.
    Meta is a copy of `d`, where data arrays are replaced with their position in the data tuple.

    Parameters
    ----------
    d: dict
        A result of ``to_dict`` method
    squeeze: bool
        If True, and there is a single data tensor, it is unpacked from the tuple

    Return:
    data, meta
    """
    data = []  # list of datas
    meta = _split_data_and_meta(d, data)
    data = data[0] if squeeze and len(data) == 1 else tuple(data)
    return data, meta



def _split_data_and_meta(d, data):
    meta = {}
    for k in sorted(d):
        if k == "data":
            data.append(d[k])
            meta[k] = len(data) - 1
        elif isinstance(d[k], dict):
            meta[k] = _split_data_and_meta(d[k], data)
        else:
            meta[k] = d[k]
    return meta




def combine_data_and_meta(data: tuple['numpy.array' | 'torch.tensor'], meta: dict) -> dict:
    """
    Reverse :meth:`yastn.split_data_and_meta`.
    """
    d = {}
    if not isinstance(data, (list, tuple)):
        data = [data]
    for k in sorted(meta):
        if k == "data":
            d[k] = data[meta[k]]
        elif isinstance(meta[k], dict):
            d[k] = combine_data_and_meta(data, meta[k])
        else:
            d[k] = meta[k]
    return d