1# Created: 13.03.2011
2# Copyright (c) 2011-2020, Manfred Moitzi
3# License: MIT License
4from typing import TYPE_CHECKING, Iterable, List, Iterator
5from itertools import chain
6
7from ezdxf.lldxf.tags import DXFStructureError
8from ezdxf.entities import entity_linker
9
10if TYPE_CHECKING:
11    from ezdxf.eztypes import (
12        TagWriter, Drawing, DXFEntity, Tags, DXFTagStorage, DXFGraphic,
13        BlockRecord,
14    )
15
16
17class StoredSection:
18    def __init__(self, entities: List['Tags']):
19        self.entities = entities
20
21    def export_dxf(self, tagwriter: 'TagWriter'):
22        # (0, SECTION) (2, NAME) is stored in entities
23        for entity in self.entities:
24            tagwriter.write_tags(entity)
25        # ENDSEC not stored in entities !!!
26        tagwriter.write_str('  0\nENDSEC\n')
27
28
29class EntitySection:
30    """ :class:`EntitiesSection` is just a proxy for :class:`Modelspace` and
31    active :class:`Paperspace` linked together.
32    """
33
34    def __init__(self, doc: 'Drawing' = None,
35                 entities: Iterable['DXFEntity'] = None):
36        self.doc = doc
37        if entities is not None:
38            self._build(iter(entities))
39
40    def __iter__(self) -> Iterable['DXFEntity']:
41        """ Iterable for all entities of modelspace and active paperspace. """
42        layouts = self.doc.layouts
43        for entity in chain(layouts.modelspace(), layouts.active_layout()):
44            yield entity
45
46    def __len__(self) -> int:
47        """ Returns count of all entities of modelspace and active paperspace.
48        """
49        layouts = self.doc.layouts
50        return len(layouts.modelspace()) + len(layouts.active_layout())
51
52    # none public interface
53
54    def _build(self, entities: Iterator['DXFEntity']) -> None:
55        section_head: 'DXFTagStorage' = next(entities)
56        if section_head.dxftype() != 'SECTION' or section_head.base_class[
57            1] != (2, 'ENTITIES'):
58            raise DXFStructureError(
59                "Critical structure error in ENTITIES section.")
60
61        def add(entity: 'DXFGraphic'):
62            handle = entity.dxf.owner
63            # higher priority for owner handle
64            if handle == msp_layout_key:
65                paperspace = 0
66            elif handle == psp_layout_key:
67                paperspace = 1
68            else:  # paperspace flag as fallback
69                paperspace = entity.dxf.paperspace
70
71            if paperspace:
72                psp.add_entity(entity)
73            else:
74                msp.add_entity(entity)
75
76        msp: 'BlockRecord' = self.doc.block_records.get('*Model_Space')
77        psp: 'BlockRecord' = self.doc.block_records.get('*Paper_Space')
78        msp_layout_key = msp.dxf.handle
79        psp_layout_key = psp.dxf.handle
80        linked_entities = entity_linker()
81        # Don't store linked entities (VERTEX, ATTRIB, SEQEND) in entity space
82        for entity in entities:
83            if not linked_entities(entity):
84                add(entity)
85
86    def export_dxf(self, tagwriter: 'TagWriter') -> None:
87        layouts = self.doc.layouts
88        tagwriter.write_str("  0\nSECTION\n  2\nENTITIES\n")
89        # Just write *Model_Space and the active *Paper_Space into the
90        # ENTITIES section.
91        layouts.modelspace().entity_space.export_dxf(tagwriter)
92        layouts.active_layout().entity_space.export_dxf(tagwriter)
93        tagwriter.write_tag2(0, "ENDSEC")
94