1# Purpose: Grouping entities by DXF attributes or a key function.
2# Created: 03.02.2017
3# Copyright (C) 2017, Manfred Moitzi
4# License: MIT License
5from typing import Iterable, Hashable, Dict, List, TYPE_CHECKING
6
7from ezdxf.lldxf.const import DXFValueError, DXFAttributeError
8
9if TYPE_CHECKING:
10    from ezdxf.eztypes import DXFEntity, KeyFunc
11
12
13def groupby(entities: Iterable['DXFEntity'], dxfattrib: str = '', key: 'KeyFunc' = None) \
14        -> Dict[Hashable, List['DXFEntity']]:
15    """
16    Groups a sequence of DXF entities by a DXF attribute like ``'layer'``, returns a dict with `dxfattrib` values
17    as key and a list of entities matching this `dxfattrib`.
18    A `key` function can be used to combine some DXF attributes (e.g. layer and color) and should return a
19    hashable data type like a tuple of strings, integers or floats, `key` function example::
20
21        def group_key(entity: DXFEntity):
22            return entity.dxf.layer, entity.dxf.color
23
24    For not suitable DXF entities return ``None`` to exclude this entity, in this case it's not required, because
25    :func:`groupby` catches :class:`DXFAttributeError` exceptions to exclude entities, which do not provide
26    layer and/or color attributes, automatically.
27
28    Result dict for `dxfattrib` = ``'layer'`` may look like this::
29
30        {
31            '0': [ ... list of entities ],
32            'ExampleLayer1': [ ... ],
33            'ExampleLayer2': [ ... ],
34            ...
35        }
36
37    Result dict for `key` = `group_key`, which returns a ``(layer, color)`` tuple, may look like this::
38
39        {
40            ('0', 1): [ ... list of entities ],
41            ('0', 3): [ ... ],
42            ('0', 7): [ ... ],
43            ('ExampleLayer1', 1): [ ... ],
44            ('ExampleLayer1', 2): [ ... ],
45            ('ExampleLayer1', 5): [ ... ],
46            ('ExampleLayer2', 7): [ ... ],
47            ...
48        }
49
50    All entity containers (modelspace, paperspace layouts and blocks) and the :class:`~ezdxf.query.EntityQuery` object
51    have a dedicated :meth:`groupby` method.
52
53    Args:
54        entities: sequence of DXF entities to group by a DXF attribute or a `key` function
55        dxfattrib: grouping DXF attribute like ``'layer'``
56        key: key function, which accepts a :class:`DXFEntity` as argument and returns a hashable grouping key
57             or ``None`` to ignore this entity.
58
59    """
60    if all((dxfattrib, key)):
61        raise DXFValueError('Specify a dxfattrib or a key function, but not both.')
62    if dxfattrib != '':
63        key = lambda entity: entity.get_dxf_attrib(dxfattrib, None)
64    if key is None:
65        raise DXFValueError('no valid argument found, specify a dxfattrib or a key function, but not both.')
66
67    result = dict()
68    for dxf_entity in entities:
69        if not dxf_entity.is_alive:
70            continue
71        try:
72            group_key = key(dxf_entity)
73        except DXFAttributeError:  # ignore DXF entities, which do not support all query attributes
74            continue
75        if group_key is not None:
76            group = result.setdefault(group_key, [])
77            group.append(dxf_entity)
78    return result
79