Skip to content

prefect.utilities.collections

Utilities for extensions of and operations on Python collections.

AutoEnum

Bases: 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
class MyEnum(AutoEnum):
    RED = AutoEnum.auto() # equivalent to RED = 'RED'
    BLUE = AutoEnum.auto() # equivalent to BLUE = 'BLUE'
Source code in prefect/utilities/collections.py
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
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}"

auto staticmethod

Exposes enum.auto() to avoid requiring a second import to use AutoEnum

Source code in prefect/utilities/collections.py
65
66
67
68
69
70
@staticmethod
def auto():
    """
    Exposes `enum.auto()` to avoid requiring a second import to use `AutoEnum`
    """
    return auto()

StopVisiting

Bases: BaseException

A special exception used to stop recursive visits in visit_collection.

When raised, the expression is returned without modification and recursive visits in that path will end.

Source code in prefect/utilities/collections.py
217
218
219
220
221
222
223
class StopVisiting(BaseException):
    """
    A special exception used to stop recursive visits in `visit_collection`.

    When raised, the expression is returned without modification and recursive visits
    in that path will end.
    """

batched_iterable

Yield batches of a certain size from an iterable

Parameters:

Name Type Description Default
iterable Iterable

An iterable

required
size int

The batch size to return

required

Yields:

Name Type Description
tuple Tuple[T, ...]

A batch of the iterable

Source code in prefect/utilities/collections.py
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
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 Type Description Default
dct dict

The dictionary to flatten

required
_parent Tuple

The current parent for recursion

None

Returns:

Type Description
Dict[Tuple[KT, ...], Any]

A flattened dict of the same type as dct

Source code in prefect/utilities/collections.py
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
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 Type Description Default
objects Iterable

An iterable of objects

required
types Union[Type[T], Tuple[Type[T], ...]]

A type or tuple of types to extract, defaults to all objects

object

Returns:

Type Description
Union[List[T], Dict[Type[T], T]]

If a single type is given: a list of instances of that type

Union[List[T], Dict[Type[T], T]]

If a tuple of types is given: a mapping of type to a list of instances

Source code in prefect/utilities/collections.py
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
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 Type Description Default
dct dict

The dictionary to be nested. Each key should be a tuple of keys as generated by dict_to_flatdict

required

Returns A nested dict of the same type as dct

Source code in prefect/utilities/collections.py
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
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

get_from_dict

Fetch a value from a nested dictionary or list using a sequence of keys.

This function allows to fetch a value from a deeply nested structure of dictionaries and lists using either a dot-separated string or a list of keys. If a requested key does not exist, the function returns the provided default value.

Parameters:

Name Type Description Default
dct Dict

The nested dictionary or list from which to fetch the value.

required
keys Union[str, List[str]]

The sequence of keys to use for access. Can be a dot-separated string or a list of keys. List indices can be included in the sequence as either integer keys or as string indices in square brackets.

required
default Any

The default value to return if the requested key path does not exist. Defaults to None.

None

Returns:

Type Description
Any

The fetched value if the key exists, or the default value if it does not.

get_from_dict({'a': {'b': {'c': [1, 2, 3, 4]}}}, 'a.b.c[1]') 2 get_from_dict({'a': {'b': [0, {'c': [1, 2]}]}}, ['a', 'b', 1, 'c', 1]) 2 get_from_dict({'a': {'b': [0, {'c': [1, 2]}]}}, 'a.b.1.c.2', 'default') 'default'

Source code in prefect/utilities/collections.py
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
def get_from_dict(dct: Dict, keys: Union[str, List[str]], default: Any = None) -> Any:
    """
    Fetch a value from a nested dictionary or list using a sequence of keys.

    This function allows to fetch a value from a deeply nested structure
    of dictionaries and lists using either a dot-separated string or a list
    of keys. If a requested key does not exist, the function returns the
    provided default value.

    Args:
        dct: The nested dictionary or list from which to fetch the value.
        keys: The sequence of keys to use for access. Can be a
            dot-separated string or a list of keys. List indices can be included
            in the sequence as either integer keys or as string indices in square
            brackets.
        default: The default value to return if the requested key path does not
            exist. Defaults to None.

    Returns:
        The fetched value if the key exists, or the default value if it does not.

    Examples:
    >>> get_from_dict({'a': {'b': {'c': [1, 2, 3, 4]}}}, 'a.b.c[1]')
    2
    >>> get_from_dict({'a': {'b': [0, {'c': [1, 2]}]}}, ['a', 'b', 1, 'c', 1])
    2
    >>> get_from_dict({'a': {'b': [0, {'c': [1, 2]}]}}, 'a.b.1.c.2', 'default')
    'default'
    """
    if isinstance(keys, str):
        keys = keys.replace("[", ".").replace("]", "").split(".")
    try:
        for key in keys:
            try:
                # Try to cast to int to handle list indices
                key = int(key)
            except ValueError:
                # If it's not an int, use the key as-is
                # for dict lookup
                pass
            dct = dct[key]
        return dct
    except (TypeError, KeyError, IndexError):
        return default

isiterable

Return a boolean indicating if an object is iterable.

Excludes types that are iterable but typically used as singletons: - str - bytes - IO objects

Source code in prefect/utilities/collections.py
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
def isiterable(obj: Any) -> bool:
    """
    Return a boolean indicating if an object is iterable.

    Excludes types that are iterable but typically used as singletons:
    - str
    - bytes
    - IO objects
    """
    try:
        iter(obj)
    except TypeError:
        return False
    else:
        return not isinstance(obj, (str, bytes, io.IOBase))

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 Type Description Default
keys_to_remove List[Hashable]

A list of keys to remove from obj obj: The object to remove keys from.

required

Returns:

Type Description

obj without keys matching an entry in keys_to_remove if obj is a dictionary. obj if obj is not a dictionary.

Source code in prefect/utilities/collections.py
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
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

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 penalties 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 - Prefect annotations

Parameters:

Name Type Description Default
expr Any

a Python object or expression

required
visit_fn Callable[[Any], Awaitable[Any]]

an async function that will be applied to every non-collection element of expr.

required
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).

False
max_depth int

Controls the depth of recursive visitation. If set to zero, no recursion will occur. If set to a positive integer N, visitation will only descend to N layers deep. If set to any negative integer, no limit will be enforced and recursion will continue until terminal items are reached. By default, recursion is unlimited.

-1
context Optional[dict]

An optional dictionary. If passed, the context will be sent to each call to the visit_fn. The context can be mutated by each visitor and will be available for later visits to expressions at the given depth. Values will not be available "up" a level from a given expression.

The context will be automatically populated with an 'annotation' key when visiting collections within a BaseAnnotation type. This requires the caller to pass context={} and will not be activated by default.

None
remove_annotations bool

If set, annotations will be replaced by their contents. By default, annotations are preserved but their contents are visited.

False
Source code in prefect/utilities/collections.py
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
def visit_collection(
    expr,
    visit_fn: Callable[[Any], Any],
    return_data: bool = False,
    max_depth: int = -1,
    context: Optional[dict] = None,
    remove_annotations: 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 penalties 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
    - Prefect annotations

    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).
        max_depth: Controls the depth of recursive visitation. If set to zero, no
            recursion will occur. If set to a positive integer N, visitation will only
            descend to N layers deep. If set to any negative integer, no limit will be
            enforced and recursion will continue until terminal items are reached. By
            default, recursion is unlimited.
        context: An optional dictionary. If passed, the context will be sent to each
            call to the `visit_fn`. The context can be mutated by each visitor and will
            be available for later visits to expressions at the given depth. Values
            will not be available "up" a level from a given expression.

            The context will be automatically populated with an 'annotation' key when
            visiting collections within a `BaseAnnotation` type. This requires the
            caller to pass `context={}` and will not be activated by default.
        remove_annotations: If set, annotations will be replaced by their contents. By
            default, annotations are preserved but their contents are visited.
    """

    def visit_nested(expr):
        # Utility for a recursive call, preserving options and updating the depth.
        return visit_collection(
            expr,
            visit_fn=visit_fn,
            return_data=return_data,
            remove_annotations=remove_annotations,
            max_depth=max_depth - 1,
            # Copy the context on nested calls so it does not "propagate up"
            context=context.copy() if context is not None else None,
        )

    def visit_expression(expr):
        if context is not None:
            return visit_fn(expr, context)
        else:
            return visit_fn(expr)

    # Visit every expression
    try:
        result = visit_expression(expr)
    except StopVisiting:
        max_depth = 0
        result = 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

    # If we have reached the maximum depth, do not perform any recursion
    if max_depth == 0:
        return result if return_data else None

    # Get the expression type; treat iterators like lists
    typ = list if isinstance(expr, IteratorABC) and isiterable(expr) 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 isinstance(expr, BaseAnnotation):
        if context is not None:
            context["annotation"] = expr
        value = visit_nested(expr.unwrap())

        if remove_annotations:
            result = value if return_data else None
        else:
            result = expr.rewrap(value) if return_data else None

    elif typ in (list, tuple, set):
        items = [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
        items = [(visit_nested(k), visit_nested(v)) for k, v in expr.items()]
        result = typ(items) if return_data else None

    elif is_dataclass(expr) and not isinstance(expr, type):
        values = [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
        # Check for presence of attrs even if they're in the field set due to pydantic#4916
        model_fields = {
            f for f in expr.__fields_set__.union(expr.__fields__) if hasattr(expr, f)
        }
        items = [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))

            # Preserve data about which fields were explicitly set on the original model
            object.__setattr__(model_instance, "__fields_set__", expr.__fields_set__)
            result = model_instance
        else:
            result = None

    else:
        result = result if return_data else None

    return result