Skip to content

treex.TreeMeta

Source code in treeo/tree.py
class TreeMeta(ABCMeta):
    def __call__(cls, *args, **kwargs) -> "Tree":
        obj: Tree

        if _COMPACT_CONTEXT.existing_subtrees is not None:
            if len(_COMPACT_CONTEXT.existing_subtrees) <= _COMPACT_CONTEXT.tree_idx:
                raise ValueError(
                    f"Out of bounds: trying to new initialize new Tree '{cls.__name__}' after the first compact run"
                )

            obj = _COMPACT_CONTEXT.existing_subtrees[_COMPACT_CONTEXT.tree_idx]
            _COMPACT_CONTEXT.tree_idx += 1

            return obj
        else:
            obj = cls.__new__(cls)
            obj = cls.construct(obj, *args, **kwargs)

        if _COMPACT_CONTEXT.new_subtrees is not None:
            _COMPACT_CONTEXT.new_subtrees.append(obj)

        return obj

    def construct(cls, obj: T, *args, **kwargs) -> T:

        obj._field_metadata = obj._field_metadata.copy()

        # set default fields
        for field, default_factory in obj._factory_fields.items():
            setattr(obj, field, default_factory())

        for field, default_value in obj._default_field_values.items():
            setattr(obj, field, default_value)

        # reset context before __init__ and add obj as current tree
        with _CompactContext(current_tree=obj):
            obj.__init__(*args, **kwargs)

        # auto-annotations
        obj._update_local_metadata()

        return obj