|
6 | 6 | from typing import Any |
7 | 7 |
|
8 | 8 | import pytest |
9 | | -from attrs import define |
| 9 | +from attrs import define, has, frozen |
10 | 10 |
|
11 | 11 | from cattrs import Converter, override |
| 12 | +from cattrs.gen import make_dict_structure_fn |
12 | 13 | from cattrs.errors import ClassValidationError, StructureHandlerNotFoundError |
13 | 14 | from cattrs.strategies import configure_tagged_union, include_subclasses |
14 | 15 |
|
@@ -536,3 +537,60 @@ class Sub(Mid1, Mid2): |
536 | 537 | assert genconverter.structure({"_type": "Sub"}, Base) == Sub() |
537 | 538 | assert genconverter.structure({"_type": "Mid1"}, Base) == Mid1() |
538 | 539 | assert genconverter.structure({"_type": "Mid2"}, Base) == Mid2() |
| 540 | + |
| 541 | + |
| 542 | +def test_subclasses_in_struct_factory(): |
| 543 | + @frozen |
| 544 | + class A: |
| 545 | + """Base class""" |
| 546 | + |
| 547 | + @frozen |
| 548 | + class A1(A): |
| 549 | + a1: int |
| 550 | + |
| 551 | + @frozen |
| 552 | + class B: |
| 553 | + id: int |
| 554 | + b: str |
| 555 | + |
| 556 | + @frozen |
| 557 | + class Container1: |
| 558 | + id: int |
| 559 | + a: A |
| 560 | + b: B |
| 561 | + |
| 562 | + @frozen |
| 563 | + class Container2: |
| 564 | + id: int |
| 565 | + c: Container1 |
| 566 | + foo: str |
| 567 | + |
| 568 | + def struct_hook_factory(cl, converter: Converter): |
| 569 | + struct_hook = make_dict_structure_fn(cl, converter) |
| 570 | + if not cl.__subclasses__(): |
| 571 | + converter.register_structure_hook(cl, struct_hook) |
| 572 | + |
| 573 | + else: |
| 574 | + |
| 575 | + def cls_is_cl(cls, _cl=cl): |
| 576 | + return cls is _cl |
| 577 | + |
| 578 | + converter.register_structure_hook_func(cls_is_cl, struct_hook) |
| 579 | + union_strategy = partial(configure_tagged_union, tag_name="type") |
| 580 | + include_subclasses(cl, converter, union_strategy=union_strategy) |
| 581 | + |
| 582 | + return converter.get_structure_hook(cl) |
| 583 | + |
| 584 | + converter = Converter() |
| 585 | + converter.register_structure_hook_factory(has, struct_hook_factory) |
| 586 | + |
| 587 | + unstructured = { |
| 588 | + "id": 0, |
| 589 | + "c": {"id": 1, "a": {"type": "A1", "a1": 42}, "b": {"id": 2, "b": "hello"}}, |
| 590 | + "foo": "world", |
| 591 | + } |
| 592 | + res = converter.structure(unstructured, Container2) |
| 593 | + |
| 594 | + assert res == Container2( |
| 595 | + id=0, c=Container1(id=1, a=A1(a1=42), b=B(id=2, b="Hello"), foo="world") |
| 596 | + ) |
0 commit comments