`prefect.utilities.collections`

Utilities for extensions of and operations on Python collections.

`AutoEnum`

An enum class that automatically generates value from variable names.

This guards against common errors where variable names are updated but values are not.

In addition, because AutoEnums inherit from str, they are automatically JSON-serializable.

See https://docs.python.org/3/library/enum.html#using-automatic-values

Examples:

class MyEnum(AutoEnum):
    RED = AutoEnum.auto() # equivalent to RED = 'RED'
    BLUE = AutoEnum.auto() # equivalent to BLUE = 'BLUE'

Source code in prefect/utilities/collections.py

class AutoEnum(str, Enum):
    """
    An enum class that automatically generates value from variable names.

    This guards against common errors where variable names are updated but values are
    not.

    In addition, because AutoEnums inherit from `str`, they are automatically
    JSON-serializable.

    See https://docs.python.org/3/library/enum.html#using-automatic-values

    Example:
        ```python
        class MyEnum(AutoEnum):
            RED = AutoEnum.auto() # equivalent to RED = 'RED'
            BLUE = AutoEnum.auto() # equivalent to BLUE = 'BLUE'
        ```
    """

    def _generate_next_value_(name, start, count, last_values):
        return name

    @staticmethod
    def auto():
        """
        Exposes `enum.auto()` to avoid requiring a second import to use `AutoEnum`
        """
        return auto()

    def __repr__(self) -> str:
        return f"{type(self).__name__}.{self.value}"

`PartialModel`

A utility for creating a Pydantic model in several steps.

Fields may be set at initialization, via attribute assignment, or at finalization when the concrete model is returned.

Pydantic validation does not occur until finalization.

Each field can only be set once and a ValueError will be raised on assignment if a field already has a value.

Examples:

>>> class MyModel(pydantic.BaseModel):
>>>     x: int
>>>     y: str
>>>     z: float
>>>
>>> partial_model = PartialModel(MyModel, x=1)
>>> partial_model.y = "two"
>>> model = partial_model.finalize(z=3.0)

Source code in prefect/utilities/collections.py

class PartialModel(Generic[M]):
    """
    A utility for creating a Pydantic model in several steps.

    Fields may be set at initialization, via attribute assignment, or at finalization
    when the concrete model is returned.

    Pydantic validation does not occur until finalization.

    Each field can only be set once and a `ValueError` will be raised on assignment if
    a field already has a value.

    Example:
        >>> class MyModel(pydantic.BaseModel):
        >>>     x: int
        >>>     y: str
        >>>     z: float
        >>>
        >>> partial_model = PartialModel(MyModel, x=1)
        >>> partial_model.y = "two"
        >>> model = partial_model.finalize(z=3.0)
    """

    def __init__(self, __model_cls: M, **kwargs: Any) -> None:
        self.fields = kwargs
        # Set fields first to avoid issues if `fields` is also set on the `model_cls`
        # in our custom `setattr` implementation.
        self.model_cls = __model_cls

        for name in kwargs.keys():
            self.raise_if_not_in_model(name)

    def finalize(self, **kwargs: Any) -> T:
        for name in kwargs.keys():
            self.raise_if_already_set(name)
            self.raise_if_not_in_model(name)
        return self.model_cls(**self.fields, **kwargs)

    def raise_if_already_set(self, name):
        if name in self.fields:
            raise ValueError(f"Field {name!r} has already been set.")

    def raise_if_not_in_model(self, name):
        if name not in self.model_cls.__fields__:
            raise ValueError(f"Field {name!r} is not present in the model.")

    def __setattr__(self, __name: str, __value: Any) -> None:
        if __name in {"fields", "model_cls"}:
            return super().__setattr__(__name, __value)

        self.raise_if_already_set(__name)
        self.raise_if_not_in_model(__name)
        self.fields[__name] = __value

    def __repr__(self) -> str:
        dsp_fields = ", ".join(f"{key}={repr(value)}" for key, value in self.fields)
        return f"PartialModel({self.model_cls.__name__}{dsp_fields})"

`Quote`

Simple wrapper to mark an expression as a different type so it will not be coerced by Prefect. For example, if you want to return a state from a flow without having the flow assume that state.

Source code in prefect/utilities/collections.py

class Quote(Generic[T]):
    """
    Simple wrapper to mark an expression as a different type so it will not be coerced
    by Prefect. For example, if you want to return a state from a flow without having
    the flow assume that state.
    """

    def __init__(self, data: T) -> None:
        self.data = data

    def unquote(self) -> T:
        return self.data

`batched_iterable`

Yield batches of a certain size from an iterable

Parameters:

Name	Description	Default
`iterable`	An iterable `Iterable`	required
`size`	The batch size to return `int`	required

Yields:

Type	Description
`tuple`	A batch of the iterable

Source code in prefect/utilities/collections.py

def batched_iterable(iterable: Iterable[T], size: int) -> Iterator[Tuple[T, ...]]:
    """
    Yield batches of a certain size from an iterable

    Args:
        iterable (Iterable): An iterable
        size (int): The batch size to return

    Yields:
        tuple: A batch of the iterable
    """
    it = iter(iterable)
    while True:
        batch = tuple(itertools.islice(it, size))
        if not batch:
            break
        yield batch

`dict_to_flatdict`

Converts a (nested) dictionary to a flattened representation.

Each key of the flat dict will be a CompoundKey tuple containing the "chain of keys" for the corresponding value.

Parameters:

Name	Description	Default
`dct`	The dictionary to flatten `dict`	required
`_parent`	The current parent for recursion `Tuple`	`None`

Returns:

Type	Description
`Dict[Tuple[~KT, ...], Any]`	A flattened dict of the same type as dct

Source code in prefect/utilities/collections.py

def dict_to_flatdict(
    dct: Dict[KT, Union[Any, Dict[KT, Any]]], _parent: Tuple[KT, ...] = None
) -> Dict[Tuple[KT, ...], Any]:
    """Converts a (nested) dictionary to a flattened representation.

    Each key of the flat dict will be a CompoundKey tuple containing the "chain of keys"
    for the corresponding value.

    Args:
        dct (dict): The dictionary to flatten
        _parent (Tuple, optional): The current parent for recursion

    Returns:
        A flattened dict of the same type as dct
    """
    typ = cast(Type[Dict[Tuple[KT, ...], Any]], type(dct))
    items: List[Tuple[Tuple[KT, ...], Any]] = []
    parent = _parent or tuple()

    for k, v in dct.items():
        k_parent = tuple(parent + (k,))
        if isinstance(v, dict):
            items.extend(dict_to_flatdict(v, _parent=k_parent).items())
        else:
            items.append((k_parent, v))
    return typ(items)

`extract_instances`

Extract objects from a file and returns a dict of type -> instances

Parameters:

Name	Description	Default
`objects`	An iterable of objects `Iterable`	required
`types`	A type or tuple of types to extract, defaults to all objects `Union[Type[~T], Tuple[Type[~T], ...]]`	`<class 'object'>`

Returns:

Type	Description
`If a single type is given`	a list of instances of that type If a tuple of types is given: a mapping of type to a list of instances

Source code in prefect/utilities/collections.py

def extract_instances(
    objects: Iterable,
    types: Union[Type[T], Tuple[Type[T], ...]] = object,
) -> Union[List[T], Dict[Type[T], T]]:
    """
    Extract objects from a file and returns a dict of type -> instances

    Args:
        objects: An iterable of objects
        types: A type or tuple of types to extract, defaults to all objects

    Returns:
        If a single type is given: a list of instances of that type
        If a tuple of types is given: a mapping of type to a list of instances
    """
    types = ensure_iterable(types)

    # Create a mapping of type -> instance from the exec values
    ret = defaultdict(list)

    for o in objects:
        # We iterate here so that the key is the passed type rather than type(o)
        for type_ in types:
            if isinstance(o, type_):
                ret[type_].append(o)

    if len(types) == 1:
        return ret[types[0]]

    return ret

`flatdict_to_dict`

Converts a flattened dictionary back to a nested dictionary.

Parameters:

Name	Description	Default
`dct`	The dictionary to be nested. Each key should be a tuple of keys as generated by `dict_to_flatdict` `dict`	required

Returns A nested dict of the same type as dct

Source code in prefect/utilities/collections.py

def flatdict_to_dict(
    dct: Dict[Tuple[KT, ...], VT]
) -> Dict[KT, Union[VT, Dict[KT, VT]]]:
    """Converts a flattened dictionary back to a nested dictionary.

    Args:
        dct (dict): The dictionary to be nested. Each key should be a tuple of keys
            as generated by `dict_to_flatdict`

    Returns
        A nested dict of the same type as dct
    """
    typ = type(dct)
    result = cast(Dict[KT, Union[VT, Dict[KT, VT]]], typ())
    for key_tuple, value in dct.items():
        current_dict = result
        for prefix_key in key_tuple[:-1]:
            # Build nested dictionaries up for the current key tuple
            # Use `setdefault` in case the nested dict has already been created
            current_dict = current_dict.setdefault(prefix_key, typ())  # type: ignore
        # Set the value
        current_dict[key_tuple[-1]] = value

    return result

`quote`

Create a Quote object

Examples:

>>> from prefect.utilities.collections import quote
>>> x = quote(1)
>>> x.unquote()
1

Source code in prefect/utilities/collections.py

def quote(expr: T) -> Quote[T]:
    """
    Create a `Quote` object

    Examples:
        >>> from prefect.utilities.collections import quote
        >>> x = quote(1)
        >>> x.unquote()
        1
    """
    return Quote(expr)

This function visits every element of an arbitrary Python collection and applies visit_fn to each element. If return_data=True, a copy of the data structure containing the results of visit_fn is returned. Note that return_data=True may be slower due to the need to copy every object.

Parameters:

Name	Description	Default
`expr`	a Python object or expression `Any`	required
`visit_fn`	an async function that will be applied to every non-collection element of expr. `Callable[[Any], Awaitable[Any]]`	required
`return_data`	if `True`, a copy of `expr` containing data modified by `visit_fn` will be returned. This is slower than `return_data=False` (the default). `bool`	`False`

Source code in prefect/utilities/collections.py

async def visit_collection(
    expr, visit_fn: Callable[[Any], Awaitable[Any]], return_data: bool = False
):
    """
    This function visits every element of an arbitrary Python collection and
    applies `visit_fn` to each element. If `return_data=True`, a copy of the
    data structure containing the results of `visit_fn` is returned. Note that
    `return_data=True` may be slower due to the need to copy every object.

    Args:
        expr (Any): a Python object or expression
        visit_fn (Callable[[Any], Awaitable[Any]]): an async function that
            will be applied to every non-collection element of expr.
        return_data (bool): if `True`, a copy of `expr` containing data modified
            by `visit_fn` will be returned. This is slower than `return_data=False`
            (the default).
    """

    def visit_nested(expr):
        # Utility for a recursive call, preserving options.
        # Returns a `partial` for use with `gather`.
        return partial(
            visit_collection, expr, visit_fn=visit_fn, return_data=return_data
        )

    # Get the expression type; treat iterators like lists
    typ = list if isinstance(expr, IteratorABC) else type(expr)
    typ = cast(type, typ)  # mypy treats this as 'object' otherwise and complains

    # do not visit mock objects
    if isinstance(expr, Mock):
        return expr if return_data else None

    elif typ in (list, tuple, set):
        result = await gather(*[visit_nested(o) for o in expr])
        return typ(result) if return_data else None

    elif typ in (dict, OrderedDict):
        assert isinstance(expr, (dict, OrderedDict))  # typecheck assertion
        keys, values = zip(*expr.items()) if expr else ([], [])
        keys = await gather(*[visit_nested(k) for k in keys])
        values = await gather(*[visit_nested(v) for v in values])
        return typ(zip(keys, values)) if return_data else None

    elif is_dataclass(expr) and not isinstance(expr, type):
        values = await gather(
            *[visit_nested(getattr(expr, f.name)) for f in fields(expr)]
        )
        result = {field.name: value for field, value in zip(fields(expr), values)}
        return typ(**result) if return_data else None

    elif (
        # Recurse into Pydantic models but do _not_ do so for states/datadocs
        isinstance(expr, pydantic.BaseModel)
        and not isinstance(expr, prefect.orion.schemas.states.State)
        and not isinstance(expr, prefect.orion.schemas.data.DataDocument)
    ):
        values = await gather(
            *[visit_nested(getattr(expr, field)) for field in expr.__fields__]
        )
        result = {field: value for field, value in zip(expr.__fields__, values)}
        return typ(**result) if return_data else None

    else:
        result = await visit_fn(expr)
        return result if return_data else None

prefect.utilities.collections

AutoEnum

PartialModel

Quote

batched_iterable

dict_to_flatdict

extract_instances

flatdict_to_dict

quote

visit_collection async

`prefect.utilities.collections`

`AutoEnum`

`PartialModel`

`Quote`

`batched_iterable`

`dict_to_flatdict`

`extract_instances`

`flatdict_to_dict`

`quote`

`visit_collection` `async`