bec/bec_lib/tests/test_signature_serializer.py

import inspect
from typing import Literal, Union

import pytest

from bec_lib.signature_serializer import (
    deserialize_dtype,
    dict_to_signature,
    serialize_dtype,
    signature_to_dict,
)


def test_signature_serializer():
    def test_func(a, b, c=1, d=2, e: int = 3):
        pass

    params = signature_to_dict(test_func)
    assert params == [
        {"name": "a", "kind": "POSITIONAL_OR_KEYWORD", "default": "_empty", "annotation": "_empty"},
        {"name": "b", "kind": "POSITIONAL_OR_KEYWORD", "default": "_empty", "annotation": "_empty"},
        {"name": "c", "kind": "POSITIONAL_OR_KEYWORD", "default": 1, "annotation": "_empty"},
        {"name": "d", "kind": "POSITIONAL_OR_KEYWORD", "default": 2, "annotation": "_empty"},
        {"name": "e", "kind": "POSITIONAL_OR_KEYWORD", "default": 3, "annotation": "int"},
    ]

    sig = dict_to_signature(params)
    assert sig == inspect.signature(test_func)


def test_signature_serializer_with_literals():
    def test_func(a, b: Literal[1, 2, 3] = 1):
        pass

    params = signature_to_dict(test_func)
    assert params == [
        {"name": "a", "kind": "POSITIONAL_OR_KEYWORD", "default": "_empty", "annotation": "_empty"},
        {
            "name": "b",
            "kind": "POSITIONAL_OR_KEYWORD",
            "default": 1,
            "annotation": {"Literal": (1, 2, 3)},
        },
    ]

    sig = dict_to_signature(params)
    assert sig == inspect.signature(test_func)


@pytest.mark.parametrize(
    "dtype_in,dtype_out",
    [
        (int, "int"),
        (str, "str"),
        (float, "float"),
        (bool, "bool"),
        (inspect._empty, "_empty"),
        # (Literal[1, 2, 3], {"Literal": (1, 2, 3)}),
        # (Union[int, str], "int | str"),
    ],
)
def test_serialize_dtype(dtype_in, dtype_out):
    assert dtype_out == serialize_dtype(dtype_in)


@pytest.mark.parametrize(
    "dtype_in,dtype_out",
    [
        ("int", int),
        ("str", str),
        ("float", float),
        ("bool", bool),
        ("_empty", inspect._empty),
        ({"Literal": (1, 2, 3)}, Literal[1, 2, 3]),
        ("int | str", Union[int, str]),
    ],
)
def test_deserialize_dtype(dtype_in, dtype_out):
    assert dtype_out == deserialize_dtype(dtype_in)