Skip to content

Commit

Permalink
adds support for fields with local dataclass types & test
Browse files Browse the repository at this point in the history
  • Loading branch information
mishamsk committed Nov 24, 2023
1 parent 517d8d5 commit d7f5a88
Show file tree
Hide file tree
Showing 2 changed files with 23 additions and 1 deletion.
10 changes: 9 additions & 1 deletion mashumaro/core/meta/code/builder.py
Original file line number Diff line number Diff line change
Expand Up @@ -80,6 +80,8 @@
__POST_SERIALIZE__ = "__post_serialize__"
__POST_DESERIALIZE__ = "__post_deserialize__"

_LOCALS_TYPE_VAR = "__locals_type_var__"


class InternalMethodName(str):
_PREFIX = "__mashumaro_"
Expand Down Expand Up @@ -303,7 +305,9 @@ def compile(self) -> None:
else:
print(f"{type_name(self.cls)}:")
print(code)
exec(code, self.globals, self.__dict__)

localns = {**self.__dict__, _LOCALS_TYPE_VAR: self.cls}
exec(code, self.globals, localns)

def evaluate_forward_ref(
self,
Expand Down Expand Up @@ -555,6 +559,10 @@ def _unpack_method_set_value(
ftype,
resolved_type_params=self.get_field_resolved_type_params(fname),
)

if "<locals>" in field_type:
field_type = _LOCALS_TYPE_VAR

could_be_none = (
ftype in (typing.Any, type(None), None)
or is_type_var_any(self._get_real_type(fname, ftype))
Expand Down
14 changes: 14 additions & 0 deletions tests/test_common.py
Original file line number Diff line number Diff line change
Expand Up @@ -262,3 +262,17 @@ def test_kw_args_when_pos_arg_is_overridden_with_field():
assert loaded.pos2 == 2
assert loaded.pos3 == 3
assert loaded.kw1 == 4


def test_local_type():
@dataclass
class LocalType:
pass

@dataclass
class DataClassWithLocalType(DataClassDictMixin):
x: LocalType

obj = DataClassWithLocalType(LocalType())
assert obj.to_dict() == {"x": {}}
assert DataClassWithLocalType.from_dict({"x": {}}) == obj

0 comments on commit d7f5a88

Please sign in to comment.