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}"
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 v is a non-empty dict, recurse
if isinstance(v, dict) and v:
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 |
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)
remove_nested_keys
¶
Recurses a dictionary returns a copy without all keys that match an entry in
key_to_remove
. Return obj
unchanged if not a dictionary.
Parameters:
Name | Description | Default |
---|---|---|
keys_to_remove |
A list of keys to remove from obj obj: The object to remove keys List[Hashable] |
required |
Returns:
Type | Description |
---|---|
|
Source code in prefect/utilities/collections.py
def remove_nested_keys(keys_to_remove: List[Hashable], obj):
"""
Recurses a dictionary returns a copy without all keys that match an entry in
`key_to_remove`. Return `obj` unchanged if not a dictionary.
Args:
keys_to_remove: A list of keys to remove from obj obj: The object to remove keys
from.
Returns:
`obj` without keys matching an entry in `keys_to_remove` if `obj` is a
dictionary. `obj` if `obj` is not a dictionary.
"""
if not isinstance(obj, dict):
return obj
return {
key: remove_nested_keys(keys_to_remove, value)
for key, value in obj.items()
if key not in keys_to_remove
}
visit_collection
async
¶
This function visits every element of an arbitrary Python collection. If an element
is a Python collection, it will be visited recursively. If an element is not a
collection, visit_fn
will be called with the element. The return value of
visit_fn
can be used to alter the element if return_data
is set.
Note that when using return_data
a copy of each collection is created to avoid
mutating the original object. This may have significant performance penalities and
should only be used if you intend to transform the collection.
Supported types: - List - Tuple - Set - Dict (note: keys are also visited recursively) - Dataclass - Pydantic model
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 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. If an element
is a Python collection, it will be visited recursively. If an element is not a
collection, `visit_fn` will be called with the element. The return value of
`visit_fn` can be used to alter the element if `return_data` is set.
Note that when using `return_data` a copy of each collection is created to avoid
mutating the original object. This may have significant performance penalities and
should only be used if you intend to transform the collection.
Supported types:
- List
- Tuple
- Set
- Dict (note: keys are also visited recursively)
- Dataclass
- Pydantic model
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,
)
# Visit every expression
result = await visit_fn(expr)
if return_data:
# Only mutate the expression while returning data, otherwise it could be null
expr = result
# Then, visit every child of the expression recursively
# 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
# Then visit every item in the expression if it is a collection
if isinstance(expr, Mock):
# Do not attempt to recurse into mock objects
result = expr
elif typ in (list, tuple, set):
items = await gather(*[visit_nested(o) for o in expr])
result = typ(items) 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])
result = 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)]
)
items = {field.name: value for field, value in zip(fields(expr), values)}
result = typ(**items) if return_data else None
elif isinstance(expr, pydantic.BaseModel):
# NOTE: This implementation *does not* traverse private attributes
# Pydantic does not expose extras in `__fields__` so we use `__fields_set__`
# as well to get all of the relevant attributes
model_fields = expr.__fields_set__.union(expr.__fields__)
items = await gather(
*[visit_nested(getattr(expr, key)) for key in model_fields]
)
if return_data:
# Collect fields with aliases so reconstruction can use the correct field name
aliases = {
key: value.alias
for key, value in expr.__fields__.items()
if value.has_alias
}
model_instance = typ(
**{
aliases.get(key) or key: value
for key, value in zip(model_fields, items)
}
)
# Private attributes are not included in `__fields_set__` but we do not want
# to drop them from the model so we restore them after constructing a new
# model
for attr in expr.__private_attributes__:
# Use `object.__setattr__` to avoid errors on immutable models
object.__setattr__(model_instance, attr, getattr(expr, attr))
result = model_instance
else:
result = None
else:
result = result if return_data else None
return result