Spaces:
Running
Running
File size: 3,296 Bytes
017ba8a |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 |
import inspect
from lynxkite.core import ops
import enum
def test_op_decorator_no_params_no_types_default_sides():
@ops.op(env="test", name="add", view=ops.ViewType.BASIC, outputs=["result"])
def add(a, b):
return a + b
assert add.__op__.name == "add"
assert add.__op__.params == {}
assert add.__op__.inputs == {
"a": ops.Input(name="a", type=inspect._empty, side=ops.Side.LEFT),
"b": ops.Input(name="b", type=inspect._empty, side=ops.Side.LEFT),
}
assert add.__op__.outputs == {
"result": ops.Output(name="result", type=None, side=ops.Side.RIGHT)
}
assert add.__op__.view_type == ops.ViewType.BASIC
assert ops.CATALOGS["test"]["add"] == add.__op__
def test_op_decorator_custom_sides():
@ops.input_side(a=ops.Side.RIGHT, b=ops.Side.TOP)
@ops.output_side(result=ops.Side.BOTTOM)
@ops.op(env="test", name="add", view=ops.ViewType.BASIC, outputs=["result"])
def add(a, b):
return a + b
assert add.__op__.name == "add"
assert add.__op__.params == {}
assert add.__op__.inputs == {
"a": ops.Input(name="a", type=inspect._empty, side=ops.Side.RIGHT),
"b": ops.Input(name="b", type=inspect._empty, side=ops.Side.TOP),
}
assert add.__op__.outputs == {
"result": ops.Output(name="result", type=None, side=ops.Side.BOTTOM)
}
assert add.__op__.view_type == ops.ViewType.BASIC
assert ops.CATALOGS["test"]["add"] == add.__op__
def test_op_decorator_with_params_and_types_():
@ops.op(env="test", name="multiply", view=ops.ViewType.BASIC, outputs=["result"])
def multiply(a: int, b: float = 2.0, *, param: str = "param"):
return a * b
assert multiply.__op__.name == "multiply"
assert multiply.__op__.params == {
"param": ops.Parameter(name="param", default="param", type=str)
}
assert multiply.__op__.inputs == {
"a": ops.Input(name="a", type=int, side=ops.Side.LEFT),
"b": ops.Input(name="b", type=float, side=ops.Side.LEFT),
}
assert multiply.__op__.outputs == {
"result": ops.Output(name="result", type=None, side=ops.Side.RIGHT)
}
assert multiply.__op__.view_type == ops.ViewType.BASIC
assert ops.CATALOGS["test"]["multiply"] == multiply.__op__
def test_op_decorator_with_complex_types():
class Color(enum.Enum):
RED = 1
GREEN = 2
BLUE = 3
@ops.op(env="test", name="color_op", view=ops.ViewType.BASIC, outputs=["result"])
def complex_op(color: Color, color_list: list[Color], color_dict: dict[str, Color]):
return color.name
assert complex_op.__op__.name == "color_op"
assert complex_op.__op__.params == {}
assert complex_op.__op__.inputs == {
"color": ops.Input(name="color", type=Color, side=ops.Side.LEFT),
"color_list": ops.Input(
name="color_list", type=list[Color], side=ops.Side.LEFT
),
"color_dict": ops.Input(
name="color_dict", type=dict[str, Color], side=ops.Side.LEFT
),
}
assert complex_op.__op__.view_type == ops.ViewType.BASIC
assert complex_op.__op__.outputs == {
"result": ops.Output(name="result", type=None, side=ops.Side.RIGHT)
}
assert ops.CATALOGS["test"]["color_op"] == complex_op.__op__
|