1# Copyright (c) 2019-2021, Manfred Moitzi
2# License: MIT License
3from typing import Optional, Iterable, Tuple, TYPE_CHECKING, Dict, Set, List
4from contextlib import contextmanager
5from ezdxf.tools.handle import HandleGenerator
6from ezdxf.lldxf.types import is_valid_handle
7from ezdxf.entities.dxfentity import DXFEntity
8from ezdxf.audit import AuditError, Auditor
9from ezdxf.lldxf.const import DXFInternalEzdxfError
10from ezdxf.entities import factory
11
12if TYPE_CHECKING:
13    from ezdxf.eztypes import TagWriter
14
15DATABASE_EXCLUDE = {
16    'SECTION', 'ENDSEC', 'EOF', 'TABLE', 'ENDTAB', 'CLASS', 'ACDSRECORD',
17    'ACDSSCHEMA'
18}
19
20
21class EntityDB:
22    """ A simple key/entity database.
23
24    Every entity/object, except tables and sections, are represented as
25    DXFEntity or inherited types, this entities are stored in the
26    DXF document database, database-key is the `handle` as string.
27
28    """
29
30    class Trashcan:
31        """ Store handles to entities which should be deleted later. """
32
33        def __init__(self, db: 'EntityDB'):
34            self._database = db._database
35            self._handles: Set[str] = set()
36
37        def add(self, handle: str):
38            """ Put handle into trashcan to delete the entity later, this is
39            required for deleting entities while iterating the database.
40            """
41            self._handles.add(handle)
42
43        def clear(self):
44            """ Remove handles in trashcan from database and destroy entities if
45            still alive.
46            """
47            db = self._database
48            for handle in self._handles:
49                entity = db.get(handle)
50                if entity and entity.is_alive:
51                    entity.destroy()
52
53                if handle in db:
54                    del db[handle]
55
56            self._handles.clear()
57
58    def __init__(self):
59        self._database: Dict[str, DXFEntity] = {}
60        # DXF handles of entities to delete later:
61        self.handles = HandleGenerator()
62        self.locked: bool = False  # used only for debugging
63
64    def __getitem__(self, handle: str) -> DXFEntity:
65        """ Get entity by `handle`, does not filter destroyed entities nor
66        entities in the trashcan.
67        """
68        return self._database[handle]
69
70    def __setitem__(self, handle: str, entity: DXFEntity) -> None:
71        """ Set `entity` for `handle`. """
72        assert isinstance(handle, str), type(handle)
73        assert isinstance(entity, DXFEntity), type(entity)
74        assert entity.is_alive, 'Can not store destroyed entity.'
75        if self.locked:
76            raise DXFInternalEzdxfError('Locked entity database.')
77
78        if handle == '0' or not is_valid_handle(handle):
79            raise ValueError(f'Invalid handle {handle}.')
80        self._database[handle] = entity
81
82    def __delitem__(self, handle: str) -> None:
83        """ Delete entity by `handle`. Removes entity only from database, does
84        not destroy the entity.
85        """
86        if self.locked:
87            raise DXFInternalEzdxfError('Locked entity database.')
88        del self._database[handle]
89
90    def __contains__(self, handle: str) -> bool:
91        """ ``True`` if database contains `handle`. """
92        if handle is None:
93            return False
94        assert isinstance(handle, str), type(handle)
95        return handle in self._database
96
97    def __len__(self) -> int:
98        """ Count of database items. """
99        return len(self._database)
100
101    def __iter__(self) -> Iterable[str]:
102        """ Iterable of all handles, does filter destroyed entities but not
103        entities in the trashcan.
104        """
105        return self.keys()
106
107    def get(self, handle: str) -> Optional[DXFEntity]:
108        """ Returns entity for `handle` or ``None`` if no entry exist, does
109        not filter destroyed entities.
110        """
111        return self._database.get(handle)
112
113    def next_handle(self) -> str:
114        """ Returns next unique handle."""
115        while True:
116            handle = self.handles.next()
117            if handle not in self._database:
118                return handle
119
120    def keys(self) -> Iterable[str]:
121        """ Iterable of all handles, does filter destroyed entities.
122        """
123        return (handle for handle, entity in self.items())
124
125    def values(self) -> Iterable[DXFEntity]:
126        """ Iterable of all entities, does filter destroyed entities.
127        """
128        return (entity for handle, entity in self.items())
129
130    def items(self) -> Iterable[Tuple[str, DXFEntity]]:
131        """ Iterable of all (handle, entities) pairs, does filter destroyed
132        entities.
133        """
134        return (
135            (handle, entity) for handle, entity in self._database.items()
136            if entity.is_alive
137        )
138
139    def add(self, entity: DXFEntity) -> None:
140        """ Add `entity` to database, assigns a new handle to the `entity`
141        if :attr:`entity.dxf.handle` is ``None``. Adding the same entity
142        multiple times is possible and creates only a single database entry.
143
144        """
145        if entity.dxftype() in DATABASE_EXCLUDE:
146            if entity.dxf.handle is not None:
147                # Mark existing entity handle as used to avoid
148                # reassigning the same handle again.
149                self[entity.dxf.handle] = entity
150            return
151        handle: str = entity.dxf.handle
152        if handle is None:
153            handle = self.next_handle()
154            entity.update_handle(handle)
155        self[handle] = entity
156
157        # Add sub entities ATTRIB, VERTEX and SEQEND to database:
158        # Add linked MTEXT columns to database:
159        if hasattr(entity, 'add_sub_entities_to_entitydb'):
160            entity.add_sub_entities_to_entitydb(self)
161
162    def delete_entity(self, entity: DXFEntity) -> None:
163        """ Remove `entity` from database and destroy the `entity`. """
164        if entity.is_alive:
165            del self[entity.dxf.handle]
166            entity.destroy()
167
168    def discard(self, entity: DXFEntity) -> None:
169        """ Discard `entity` from database without destroying the `entity`. """
170        if entity.is_alive:
171            if hasattr(entity, 'process_sub_entities'):
172                entity.process_sub_entities(lambda e: self.discard(e))
173
174            handle = entity.dxf.handle
175            try:
176                del self._database[handle]
177                entity.dxf.handle = None
178            except KeyError:
179                pass
180
181    def duplicate_entity(self, entity: DXFEntity) -> DXFEntity:
182        """ Duplicates `entity` and its sub entities (VERTEX, ATTRIB, SEQEND)
183        and store them with new handles in the entity database.
184        Graphical entities have to be added to a layout by
185        :meth:`~ezdxf.layouts.BaseLayout.add_entity`.
186
187        To import DXF entities from another drawing use the
188        :class:`~ezdxf.addons.importer.Importer` add-on.
189
190        A new owner handle will be set by adding the duplicated entity to a
191        layout.
192
193        """
194        new_entity: DXFEntity = entity.copy()
195        new_entity.dxf.handle = self.next_handle()
196        factory.bind(new_entity, entity.doc)
197        return new_entity
198
199    def audit(self, auditor: 'Auditor'):
200        """ Restore database integrity:
201
202        - restore database entries with modified handles (key != entity.dxf.handle)
203        - remove entities with invalid handles
204        - empty trashcan - destroy all entities in the trashcan
205        - removes destroyed database entries (purge)
206
207        """
208        assert self.locked is False, 'Database is locked!'
209        add_entities = []
210
211        with self.trashcan() as trash:
212            for handle, entity in self.items():
213                # Destroyed entities are already filtered!
214                if not is_valid_handle(handle):
215                    auditor.fixed_error(
216                        code=AuditError.INVALID_ENTITY_HANDLE,
217                        message=f'Removed entity {entity.dxftype()} with invalid '
218                                f'handle "{handle}" from entity database.',
219                    )
220                    trash.add(handle)
221                if handle != entity.dxf.get('handle'):
222                    # database handle != stored entity handle
223                    # prevent entity from being destroyed:
224                    self._database[handle] = None
225                    trash.add(handle)
226                    add_entities.append(entity)
227
228        # Remove all destroyed entities from database:
229        self.purge()
230
231        for entity in add_entities:
232            handle = entity.dxf.get('handle')
233            if handle is None:
234                auditor.fixed_error(
235                    code=AuditError.INVALID_ENTITY_HANDLE,
236                    message=f'Removed entity {entity.dxftype()} without handle '
237                            f'from entity database.',
238                )
239                continue
240            if not is_valid_handle(handle) or handle == '0':
241                auditor.fixed_error(
242                    code=AuditError.INVALID_ENTITY_HANDLE,
243                    message=f'Removed entity {entity.dxftype()} with invalid '
244                            f'handle "{handle}" from entity database.',
245                )
246                continue
247            self[handle] = entity
248
249    def new_trashcan(self) -> 'EntityDB.Trashcan':
250        """ Returns a new trashcan, empty trashcan manually by: :
251        func:`Trashcan.clear()`.
252        """
253        return EntityDB.Trashcan(self)
254
255    @contextmanager
256    def trashcan(self) -> 'EntityDB.Trashcan':
257        """ Returns a new trashcan in context manager mode, trashcan will be
258        emptied when leaving context.
259        """
260        trashcan_ = self.new_trashcan()
261        yield trashcan_
262        # try ... finally is not required, in case of an exception the database
263        # is maybe already in an unreliable state.
264        trashcan_.clear()
265
266    def purge(self) -> None:
267        """ Remove all destroyed entities from database, but does not empty the
268        trashcan.
269        """
270        # Important: operate on underlying data structure:
271        db = self._database
272        dead_handles = [
273            handle for handle, entity in db.items()
274            if not entity.is_alive
275        ]
276        for handle in dead_handles:
277            del db[handle]
278
279    def dxf_types_in_use(self) -> Set[str]:
280        return set(entity.dxftype() for entity in self.values())
281
282
283class EntitySpace:
284    """
285    An :class:`EntitySpace` is a collection of :class:`~ezdxf.entities.DXFEntity`
286    objects, that stores only  references to :class:`DXFEntity` objects.
287
288    The :class:`~ezdxf.layouts.Modelspace`, any :class:`~ezdxf.layouts.Paperspace`
289    layout and :class:`~ezdxf.layouts.BlockLayout` objects have an
290    :class:`EntitySpace` container to store their entities.
291
292    """
293
294    def __init__(self, entities: Iterable[DXFEntity] = None):
295        self.entities: List[DXFEntity] = list(e for e in entities if e.is_alive) \
296            if entities else []
297
298    def __iter__(self) -> Iterable[DXFEntity]:
299        """ Iterable of all entities, filters destroyed entities. """
300        return (e for e in self.entities if e.is_alive)
301
302    def __getitem__(self, index) -> DXFEntity:
303        """ Get entity at index `item`
304
305        :class:`EntitySpace` has a standard Python list like interface,
306        therefore `index` can be any valid list indexing or slicing term, like
307        a single index ``layout[-1]`` to get the last entity, or an index slice
308        ``layout[:10]`` to get the first 10 or less entities as
309        ``List[DXFEntity]``. Does not filter destroyed entities.
310
311        """
312        return self.entities[index]
313
314    def __len__(self) -> int:
315        """ Count of entities including destroyed entities. """
316        return len(self.entities)
317
318    def has_handle(self, handle: str) -> bool:
319        """ ``True`` if `handle` is present, does filter destroyed entities. """
320        assert isinstance(handle, str), type(handle)
321        return any(e.dxf.handle == handle for e in self)
322
323    def purge(self):
324        """ Remove all destroyed entities from entity space. """
325        self.entities = list(self)
326
327    def add(self, entity: DXFEntity) -> None:
328        """ Add `entity`. """
329        assert isinstance(entity, DXFEntity), type(entity)
330        assert entity.is_alive, 'Can not store destroyed entities'
331        self.entities.append(entity)
332
333    def extend(self, entities: Iterable[DXFEntity]) -> None:
334        """ Add multiple `entities`."""
335        for entity in entities:
336            self.add(entity)
337
338    def export_dxf(self, tagwriter: 'TagWriter') -> None:
339        """ Export all entities into DXF file by `tagwriter`.
340
341        (internal API)
342        """
343        for entity in iter(self):
344            entity.export_dxf(tagwriter)
345
346    def remove(self, entity: DXFEntity) -> None:
347        """ Remove `entity`. """
348        self.entities.remove(entity)
349
350    def clear(self) -> None:
351        """ Remove all entities. """
352        # Do not destroy entities!
353        self.entities = list()
354