diff --git a/transactron/utils/transactron_helpers.py b/transactron/utils/transactron_helpers.py index 1d41178..1fae882 100644 --- a/transactron/utils/transactron_helpers.py +++ b/transactron/utils/transactron_helpers.py @@ -19,6 +19,8 @@ "mock_def_helper", "get_src_loc", "from_method_layout", + "make_layout", + "extend_layout", ] T = TypeVar("T") @@ -147,6 +149,10 @@ def make_layout(*fields: LayoutListField) -> StructLayout: return from_method_layout(fields) +def extend_layout(layout: StructLayout, *fields: LayoutListField) -> StructLayout: + return StructLayout(layout.members | from_method_layout(fields).members) + + def from_method_layout(layout: MethodLayout) -> StructLayout: if isinstance(layout, StructLayout): return layout