diff --git a/boa3/internal/model/type/collection/icollection.py b/boa3/internal/model/type/collection/icollection.py index dae440eb..515753a9 100644 --- a/boa3/internal/model/type/collection/icollection.py +++ b/boa3/internal/model/type/collection/icollection.py @@ -2,6 +2,7 @@ from collections.abc import Iterable from typing import Any, Self +from boa3.internal.model.expression import IExpression from boa3.internal.model.type.annotation.uniontype import UnionType from boa3.internal.model.type.classes.pythonclass import PythonClass from boa3.internal.model.type.itype import IType @@ -73,7 +74,11 @@ def get_types(cls, value: Any) -> set[IType]: if not isinstance(value, Iterable): value = {value} - types: set[IType] = {val if isinstance(val, IType) else Type.get_type(val) for val in value} + types: set[IType] = { + val if isinstance(val, IType) + else Type.get_type(val.type if isinstance(val, IExpression) else val) + for val in value + } return cls.filter_types(types) def get_item_type(self, index: tuple): diff --git a/boa3_test/test_sc/class_test/ReturnDictWithClassAttributes.py b/boa3_test/test_sc/class_test/ReturnDictWithClassAttributes.py new file mode 100644 index 00000000..837f327b --- /dev/null +++ b/boa3_test/test_sc/class_test/ReturnDictWithClassAttributes.py @@ -0,0 +1,53 @@ +from typing import Any + +from boa3.builtin.compile_time import public + + +class Example: + def __init__(self, shape: str, color: str, background: str, size: str): + self.shape = shape + self.color = color + self.background = background + self.size = size + + def test_value(self) -> dict[str, str]: + return { + 'shape': self.shape, + 'color': self.color, + 'background': self.background, + 'size': self.size + } + + def test_keys(self) -> dict[str, str]: + return { + self.shape: 'shape', + self.color: 'color', + self.background: 'background', + self.size: 'size' + } + + def test_pair(self) -> dict[str, str]: + return { + self.shape: self.shape, + self.color: self.color, + self.background: self.background, + self.size: self.size + } + + +@public +def test_only_values() -> Any: + example = Example('Rectangle', 'Blue', 'Black', 'Small') + return example.test_value() + + +@public +def test_only_keys() -> Any: + example = Example('Rectangle', 'Blue', 'Black', 'Small') + return example.test_keys() + + +@public +def test_pair() -> Any: + example = Example('Rectangle', 'Blue', 'Black', 'Small') + return example.test_pair() diff --git a/boa3_test/tests/compiler_tests/test_class.py b/boa3_test/tests/compiler_tests/test_class.py index 5bd09ff5..34a52047 100644 --- a/boa3_test/tests/compiler_tests/test_class.py +++ b/boa3_test/tests/compiler_tests/test_class.py @@ -459,3 +459,33 @@ async def test_class_property_and_parameter_with_same_name(self): result, _ = await self.call('main', ['unit test'], return_type=str) self.assertEqual('unit test', result) + + async def test_return_dict_with_class_attributes(self): + await self.set_up_contract('ReturnDictWithClassAttributes.py') + + expected_result = { + 'shape': 'Rectangle', + 'color': 'Blue', + 'background': 'Black', + 'size': 'Small' + } + result, _ = await self.call('test_only_values', [], return_type=dict[str,str]) + self.assertEqual(expected_result, result) + + expected_result = { + 'Rectangle': 'shape', + 'Blue': 'color', + 'Black': 'background', + 'Small': 'size' + } + result, _ = await self.call('test_only_keys', [], return_type=dict[str,str]) + self.assertEqual(expected_result, result) + + expected_result = { + 'Rectangle': 'Rectangle', + 'Blue': 'Blue', + 'Black': 'Black', + 'Small': 'Small' + } + result, _ = await self.call('test_pair', [], return_type=dict[str,str]) + self.assertEqual(expected_result, result)