treex.Compact

Source code in treeo/mixins.py

class Compact:
    _field_metadata: tp.Dict[str, types.FieldMetadata]
    _subtrees: tp.Optional[tp.Tuple[str, ...]]

    @property
    def first_run(self) -> bool:
        """
        Returns:
            `True` if its currently the first run of a `compact` method.
        """
        if tree_m._COMPACT_CONTEXT.current_tree is not self:
            raise RuntimeError(
                f"Object '{type(self).__name__}' is not the current tree, found '{type(tree_m._COMPACT_CONTEXT.current_tree).__name__}', did you forget the @compact decorator?"
            )

        return self._subtrees is None

    # NOTE: it feels like `get_field` could be safely used in non-`compact` methods, maybe
    # the various checks done to verify that this method is used inside `compact` could be removed.
    def get_field(
        self,
        field_name: str,
        initializer: tp.Callable[[], A],
    ) -> A:
        """
        A method that gets a field with the given name if exists, otherwise it initializes it and returns it.

        Currently the follow restrictions apply:

        * The field must be declared in the class definition.
        * The method can only be called inside a `compact` context.

        Arguments:
            field_name: The name of the field to get.
            initializer: The function to initialize the field if it does not exist.

        Returns:
            The field value.
        """
        value: A

        if field_name not in self._field_metadata:
            raise ValueError(f"Metadata for field '{field_name}' does not exist.")

        if field_name in vars(self):
            value = getattr(self, field_name)
        else:
            if tree_m._COMPACT_CONTEXT.in_compact and not self.first_run:
                raise RuntimeError(
                    f"Trying to initialize field '{field_name}' after the first run of `compact`."
                )

            value = initializer()
            setattr(self, field_name, value)

        return value

first_run: bool property readonly

Returns:

Type	Description
bool	True if its currently the first run of a compact method.

get_field(self, field_name, initializer)

A method that gets a field with the given name if exists, otherwise it initializes it and returns it.

Currently the follow restrictions apply:

• The field must be declared in the class definition.
• The method can only be called inside a compact context.

Parameters:

Name	Type	Description	Default
field_name	str	The name of the field to get.	required
initializer	Callable[[], ~A]	The function to initialize the field if it does not exist.	required

Returns:

Type	Description
~A	The field value.

Source code in treeo/mixins.py

def get_field(
    self,
    field_name: str,
    initializer: tp.Callable[[], A],
) -> A:
    """
    A method that gets a field with the given name if exists, otherwise it initializes it and returns it.

    Currently the follow restrictions apply:

    * The field must be declared in the class definition.
    * The method can only be called inside a `compact` context.

    Arguments:
        field_name: The name of the field to get.
        initializer: The function to initialize the field if it does not exist.

    Returns:
        The field value.
    """
    value: A

    if field_name not in self._field_metadata:
        raise ValueError(f"Metadata for field '{field_name}' does not exist.")

    if field_name in vars(self):
        value = getattr(self, field_name)
    else:
        if tree_m._COMPACT_CONTEXT.in_compact and not self.first_run:
            raise RuntimeError(
                f"Trying to initialize field '{field_name}' after the first run of `compact`."
            )

        value = initializer()
        setattr(self, field_name, value)

    return value